Refactor naming and behavior of bulk load methods in TransactionManager (#918)

* Refactor naming and behavior of bulk load methods in TransactionManager

The contract of loadByKeys(Iterable<VKey>) specifies that the method will
throw a NoSuchElementException if any of the specified keys don't exist.
We don't do that before this PR, but now we do.

Existing calls (when necessary) were converted to the new load*
methods, which have the same behavior as the previous methods.

Existing methods were also renamed to be more clear -- see b/176239831
for more details and discussion.

core/src/main/java/google/registry/batch/DeleteContactsAndHostsAction.java

@@ -345,7 +345,7 @@ public class DeleteContactsAndHostsAction implements Runnable {
       String resourceClientId = resource.getPersistedCurrentSponsorClientId();
       if (resource instanceof HostResource && ((HostResource) resource).isSubordinate()) {
         resourceClientId =
-            tm().load(((HostResource) resource).getSuperordinateDomain())
+            tm().loadByKey(((HostResource) resource).getSuperordinateDomain())
                 .cloneProjectedAtTime(now)
                 .getCurrentSponsorClientId();
       }
@@ -465,7 +465,7 @@ public class DeleteContactsAndHostsAction implements Runnable {
         if (host.isSubordinate()) {
           dnsQueue.addHostRefreshTask(host.getHostName());
           tm().put(
-                  tm().load(host.getSuperordinateDomain())
+                  tm().loadByKey(host.getSuperordinateDomain())
                       .asBuilder()
                       .removeSubordinateHost(host.getHostName())
                       .build());

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

@@ -139,7 +139,7 @@ public final class ResourceFlowUtils {
       Class<R> clazz, String targetId, DateTime now, String clientId) throws EppException {
     VKey<R> key = loadAndGetKey(clazz, targetId, now);
     if (key != null) {
-      R resource = tm().load(key);
+      R resource = tm().loadByKey(key);
       // These are similar exceptions, but we can track them internally as log-based metrics.
       if (Objects.equals(clientId, resource.getPersistedCurrentSponsorClientId())) {
         throw new ResourceAlreadyExistsForThisClientException(targetId);

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

@@ -225,7 +225,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
-          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
+          BillingEvent.OneTime oneTime = tm().loadByKey(gracePeriod.getOneTimeBillingEvent());
           newExpirationTime = newExpirationTime.minusYears(oneTime.getPeriodYears());
         } else if (gracePeriod.getRecurringBillingEvent() != null) {
           // Take 1 year off the registration if in the autorenew grace period (no need to load the
@@ -372,12 +372,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
       DateTime autoRenewTime =
-          tm().load(checkNotNull(gracePeriod.getRecurringBillingEvent()))
+          tm().loadByKey(checkNotNull(gracePeriod.getRecurringBillingEvent()))
               .getRecurrenceTimeOfYear()
               .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
-    return tm().load(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
+    return tm().loadByKey(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
   }
 
   @Nullable

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

@@ -517,7 +517,7 @@ public class DomainFlowUtils {
    */
   public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
-        tm().maybeLoad(domain.getAutorenewPollMessage());
+        tm().loadByKeyIfPresent(domain.getAutorenewPollMessage());
 
     // Construct an updated autorenew poll message. If the autorenew poll message no longer exists,
     // create a new one at the same id. This can happen if a transfer was requested on a domain
@@ -542,7 +542,7 @@ public class DomainFlowUtils {
       ofy().save().entity(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = tm().load(domain.getAutorenewBillingEvent());
+    Recurring recurring = tm().loadByKey(domain.getAutorenewBillingEvent());
     ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 
@@ -1022,7 +1022,7 @@ public class DomainFlowUtils {
     for (DesignatedContact contact : contacts) {
       builder.add(
           ForeignKeyedDesignatedContact.create(
-              contact.getType(), tm().load(contact.getContactKey()).getContactId()));
+              contact.getType(), tm().loadByKey(contact.getContactKey()).getContactId()));
     }
     return builder.build();
   }

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

@@ -101,8 +101,8 @@ public final class DomainInfoFlow implements Flow {
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());
     // Prefetch all referenced resources. Calling values() blocks until loading is done.
-    tm().load(domain.getNameservers());
-    tm().load(domain.getReferencedContacts());
+    tm().loadByKeys(domain.getNameservers());
+    tm().loadByKeys(domain.getReferencedContacts());
     // Registrars can only see a few fields on unauthorized domains.
     // This is a policy decision that is left up to us by the rfcs.
     DomainInfoData.Builder infoBuilder =
@@ -110,7 +110,7 @@ public final class DomainInfoFlow implements Flow {
             .setFullyQualifiedDomainName(domain.getDomainName())
             .setRepoId(domain.getRepoId())
             .setCurrentSponsorClientId(domain.getCurrentSponsorClientId())
-            .setRegistrant(tm().load(domain.getRegistrant()).getContactId());
+            .setRegistrant(tm().loadByKey(domain.getRegistrant()).getContactId());
     // If authInfo is non-null, then the caller is authorized to see the full information since we
     // will have already verified the authInfo is valid.
     if (clientId.equals(domain.getCurrentSponsorClientId()) || authInfo.isPresent()) {

core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java

@@ -153,7 +153,7 @@ public class AllocationTokenFlowUtils {
       throw new InvalidAllocationTokenException();
     }
     Optional<AllocationToken> maybeTokenEntity =
-        tm().maybeLoad(VKey.create(AllocationToken.class, token));
+        tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token));
     if (maybeTokenEntity.isEmpty()) {
       throw new InvalidAllocationTokenException();
     }

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

@@ -96,7 +96,7 @@ public final class HostDeleteFlow implements TransactionalFlow {
       // the client id, needs to be read off of it.
       EppResource owningResource =
           existingHost.isSubordinate()
-              ? tm().load(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
+              ? tm().loadByKey(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
               : existingHost;
       verifyResourceOwnership(clientId, owningResource);
     }

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

@@ -77,7 +77,7 @@ public final class HostInfoFlow implements Flow {
     // there is no superordinate domain, the host's own values for these fields will be correct.
     if (host.isSubordinate()) {
       DomainBase superordinateDomain =
-          tm().load(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
+          tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
       hostInfoDataBuilder
           .setCurrentSponsorClientId(superordinateDomain.getCurrentSponsorClientId())
           .setLastTransferTime(host.computeLastTransferTime(superordinateDomain));

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

@@ -139,7 +139,7 @@ public final class HostUpdateFlow implements TransactionalFlow {
     String newHostName = firstNonNull(suppliedNewHostName, oldHostName);
     DomainBase oldSuperordinateDomain =
         existingHost.isSubordinate()
-            ? tm().load(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
+            ? tm().loadByKey(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
             : null;
     // Note that lookupSuperordinateDomain calls cloneProjectedAtTime on the domain for us.
     Optional<DomainBase> newSuperordinateDomain =
@@ -286,7 +286,7 @@ public final class HostUpdateFlow implements TransactionalFlow {
         && Objects.equals(
             existingHost.getSuperordinateDomain(), newHost.getSuperordinateDomain())) {
       tm().put(
-              tm().load(existingHost.getSuperordinateDomain())
+              tm().loadByKey(existingHost.getSuperordinateDomain())
                   .asBuilder()
                   .removeSubordinateHost(existingHost.getHostName())
                   .addSubordinateHost(newHost.getHostName())
@@ -295,14 +295,14 @@ public final class HostUpdateFlow implements TransactionalFlow {
     }
     if (existingHost.isSubordinate()) {
       tm().put(
-              tm().load(existingHost.getSuperordinateDomain())
+              tm().loadByKey(existingHost.getSuperordinateDomain())
                   .asBuilder()
                   .removeSubordinateHost(existingHost.getHostName())
                   .build());
     }
     if (newHost.isSubordinate()) {
       tm().put(
-              tm().load(newHost.getSuperordinateDomain())
+              tm().loadByKey(newHost.getSuperordinateDomain())
                   .asBuilder()
                   .addSubordinateHost(newHost.getHostName())
                   .build());

core/src/main/java/google/registry/model/EppResource.java

@@ -360,13 +360,13 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
 
         @Override
         public EppResource load(VKey<? extends EppResource> key) {
-          return tm().doTransactionless(() -> tm().load(key));
+          return tm().doTransactionless(() -> tm().loadByKey(key));
         }
 
         @Override
         public Map<VKey<? extends EppResource>, EppResource> loadAll(
             Iterable<? extends VKey<? extends EppResource>> keys) {
-          return tm().doTransactionless(() -> tm().load(keys));
+          return tm().doTransactionless(() -> tm().loadByKeys(keys));
         }
       };
 
@@ -406,7 +406,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   public static ImmutableMap<VKey<? extends EppResource>, EppResource> loadCached(
       Iterable<VKey<? extends EppResource>> keys) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
-      return tm().load(keys);
+      return tm().loadByKeys(keys);
     }
     try {
       return cacheEppResources.getAll(keys);
@@ -423,7 +423,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
    */
   public static <T extends EppResource> T loadCached(VKey<T> key) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
-      return tm().load(key);
+      return tm().loadByKey(key);
     }
     try {
       // Safe to cast because loading a Key<T> returns an entity of type T.

core/src/main/java/google/registry/model/EppResourceUtils.java

@@ -144,7 +144,7 @@ public final class EppResourceUtils {
     T resource =
         useCache
             ? EppResource.loadCached(fki.getResourceKey())
-            : transactIfJpaTm(() -> tm().maybeLoad(fki.getResourceKey()).orElse(null));
+            : transactIfJpaTm(() -> tm().loadByKeyIfPresent(fki.getResourceKey()).orElse(null));
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return Optional.empty();
     }

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -16,6 +16,7 @@ package google.registry.model.ofy;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -23,6 +24,8 @@ import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
@@ -177,26 +180,13 @@ public class DatastoreTransactionManager implements TransactionManager {
   // VKey instead of by ofy Key.  But ideally, there should be one set of TransactionManager
   // interface tests that are applied to both the datastore and SQL implementations.
   @Override
-  public <T> Optional<T> maybeLoad(VKey<T> key) {
+  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     return Optional.ofNullable(loadNullable(key));
   }
 
   @Override
-  public <T> T load(VKey<T> key) {
-    T result = loadNullable(key);
-    if (result == null) {
-      throw new NoSuchElementException(key.toString());
-    }
-    return result;
-  }
-
-  @Override
-  public <T> T load(T entity) {
-    return ofy().load().entity(entity).now();
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys) {
     // Keep track of the Key -> VKey mapping so we can translate them back.
     ImmutableMap<Key<T>, VKey<? extends T>> keyMap =
         StreamSupport.stream(keys.spliterator(), false)
@@ -211,13 +201,51 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
-    return ImmutableList.copyOf(getOfy().load().type(clazz));
+  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
+    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
-    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
+  public <T> T loadByKey(VKey<T> key) {
+    T result = loadNullable(key);
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
+      Iterable<? extends VKey<? extends T>> keys) {
+    ImmutableMap<VKey<? extends T>, T> result = loadByKeysIfPresent(keys);
+    ImmutableSet<? extends VKey<? extends T>> missingKeys =
+        Streams.stream(keys).filter(k -> !result.containsKey(k)).collect(toImmutableSet());
+    if (!missingKeys.isEmpty()) {
+      // Ofy ignores nonexistent keys but the method contract specifies to throw if nonexistent
+      throw new NoSuchElementException(
+          String.format("Failed to load nonexistent entities for keys: %s", missingKeys));
+    }
+    return result;
+  }
+
+  @Override
+  public <T> T loadByEntity(T entity) {
+    return ofy().load().entity(entity).now();
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
+    ImmutableList<T> result = loadByEntitiesIfPresent(entities);
+    if (result.size() != Iterables.size(entities)) {
+      throw new NoSuchElementException(
+          String.format("Attempted to load entities, some of which are missing: %s", entities));
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
+    return ImmutableList.copyOf(getOfy().load().type(clazz));
   }
 
   @Override

core/src/main/java/google/registry/model/rde/RdeRevision.java

@@ -97,7 +97,7 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
     Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, id);
     Optional<RdeRevision> revisionOptional =
-        tm().maybeLoad(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
     return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
   }
 
@@ -117,7 +117,7 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
     Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, triplet);
     Optional<RdeRevision> revisionOptional =
-        tm().maybeLoad(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
     if (revision == 0) {
       revisionOptional.ifPresent(
           rdeRevision -> {

core/src/main/java/google/registry/model/registrar/Registrar.java

@@ -815,7 +815,8 @@ public class Registrar extends ImmutableObject
               .map(Registry::createVKey)
               .collect(toImmutableSet());
       Set<VKey<Registry>> missingTldKeys =
-          Sets.difference(newTldKeys, transactIfJpaTm(() -> tm().load(newTldKeys)).keySet());
+          Sets.difference(
+              newTldKeys, transactIfJpaTm(() -> tm().loadByKeysIfPresent(newTldKeys)).keySet());
       checkArgument(missingTldKeys.isEmpty(), "Trying to set nonexisting TLDs: %s", missingTldKeys);
       getInstance().allowedTlds = ImmutableSortedSet.copyOf(allowedTlds);
       return this;
@@ -983,7 +984,7 @@ public class Registrar extends ImmutableObject
   public static Iterable<Registrar> loadAll() {
     return tm().isOfy()
         ? ImmutableList.copyOf(ofy().load().type(Registrar.class).ancestor(getCrossTldKey()))
-        : tm().transact(() -> tm().loadAll(Registrar.class));
+        : tm().transact(() -> tm().loadAllOf(Registrar.class));
   }
 
   /** Loads all registrar entities using an in-memory cache. */
@@ -994,7 +995,7 @@ public class Registrar extends ImmutableObject
   /** Loads and returns a registrar entity by its client id directly from Datastore. */
   public static Optional<Registrar> loadByClientId(String clientId) {
     checkArgument(!Strings.isNullOrEmpty(clientId), "clientId must be specified");
-    return transactIfJpaTm(() -> tm().maybeLoad(createVKey(clientId)));
+    return transactIfJpaTm(() -> tm().loadByKeyIfPresent(createVKey(clientId)));
   }
 
   /**

core/src/main/java/google/registry/model/registry/Registries.java

@@ -69,7 +69,7 @@ public final class Registries {
                                   .stream()
                                   .map(Key::getName)
                                   .collect(toImmutableSet())
-                              : tm().loadAll(Registry.class).stream()
+                              : tm().loadAllOf(Registry.class).stream()
                                   .map(Registry::getTldStr)
                                   .collect(toImmutableSet());
                       return Registry.getAll(tlds).stream()

core/src/main/java/google/registry/model/registry/Registry.java

@@ -267,7 +267,7 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
                 public Optional<Registry> load(final String tld) {
                   // Enter a transaction-less context briefly; we don't want to enroll every TLD in
                   // a transaction that might be wrapping this call.
-                  return tm().doTransactionless(() -> tm().maybeLoad(createVKey(tld)));
+                  return tm().doTransactionless(() -> tm().loadByKeyIfPresent(createVKey(tld)));
                 }
 
                 @Override
@@ -275,7 +275,7 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
                   ImmutableMap<String, VKey<Registry>> keysMap =
                       toMap(ImmutableSet.copyOf(tlds), Registry::createVKey);
                   Map<VKey<? extends Registry>, Registry> entities =
-                      tm().doTransactionless(() -> tm().load(keysMap.values()));
+                      tm().doTransactionless(() -> tm().loadByKeys(keysMap.values()));
                   return Maps.transformEntries(
                       keysMap, (k, v) -> Optional.ofNullable(entities.getOrDefault(v, null)));
                 }

core/src/main/java/google/registry/model/registry/label/ReservedListDualWriteDao.java

@@ -62,7 +62,7 @@ public class ReservedListDualWriteDao {
   public static Optional<ReservedList> getLatestRevision(String reservedListName) {
     Optional<ReservedList> maybeDatastoreList =
         ofyTm()
-            .maybeLoad(
+            .loadByKeyIfPresent(
                 VKey.createOfy(
                     ReservedList.class,
                     Key.create(getCrossTldKey(), ReservedList.class, reservedListName)));

core/src/main/java/google/registry/model/server/ServerSecret.java

@@ -77,7 +77,7 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
               // transactionally create a new ServerSecret (once per app setup) if necessary.
               // return the ofy() result during Datastore-primary phase
               ServerSecret secret =
-                  ofyTm().maybeLoad(key).orElseGet(() -> create(UUID.randomUUID()));
+                  ofyTm().loadByKeyIfPresent(key).orElseGet(() -> create(UUID.randomUUID()));
               // During a dual-write period, write it to both Datastore and SQL
               // even if we didn't have to retrieve it from the DB
               ofyTm().transact(() -> ofyTm().putWithoutBackup(secret));

core/src/main/java/google/registry/model/tmch/TmchCrl.java

@@ -56,7 +56,7 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
         VKey.create(
             TmchCrl.class, SINGLETON_ID, Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID));
     // return the ofy() result during Datastore-primary phase
-    return ofyTm().transact(() -> ofyTm().maybeLoad(key));
+    return ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(key));
   }
 
   /**

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -40,7 +40,7 @@ import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
-import java.util.Map.Entry;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -356,33 +356,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public <T> Optional<T> maybeLoad(VKey<T> key) {
+  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
     return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
   }
 
   @Override
-  public <T> T load(VKey<T> key) {
-    checkArgumentNotNull(key, "key must be specified");
-    assertInTransaction();
-    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
-    if (result == null) {
-      throw new NoSuchElementException(key.toString());
-    }
-    return result;
-  }
-
-  @Override
-  public <T> T load(T entity) {
-    checkArgumentNotNull(entity, "entity must be specified");
-    assertInTransaction();
-    return (T)
-        load(VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");
     assertInTransaction();
     return StreamSupport.stream(keys.spliterator(), false)
@@ -393,11 +375,58 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
                 new SimpleEntry<VKey<? extends T>, T>(
                     key, getEntityManager().find(key.getKind(), key.getSqlKey())))
         .filter(entry -> entry.getValue() != null)
-        .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+        .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
+  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
+    return Streams.stream(entities)
+        .filter(this::exists)
+        .map(this::loadByEntity)
+        .collect(toImmutableList());
+  }
+
+  @Override
+  public <T> T loadByKey(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
+      Iterable<? extends VKey<? extends T>> keys) {
+    ImmutableMap<VKey<? extends T>, T> existing = loadByKeysIfPresent(keys);
+    ImmutableSet<? extends VKey<? extends T>> missingKeys =
+        Streams.stream(keys).filter(k -> !existing.containsKey(k)).collect(toImmutableSet());
+    if (!missingKeys.isEmpty()) {
+      throw new NoSuchElementException(
+          String.format(
+              "Expected to find the following VKeys but they were missing: %s.", missingKeys));
+    }
+    return existing;
+  }
+
+  @Override
+  public <T> T loadByEntity(T entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    return (T)
+        loadByKey(
+            VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
+    return Streams.stream(entities).map(this::loadByEntity).collect(toImmutableList());
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
     return ImmutableList.copyOf(
@@ -408,11 +437,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             .getResultList());
   }
 
-  @Override
-  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
-    return Streams.stream(entities).map(this::load).collect(toImmutableList());
-  }
-
   private int internalDelete(VKey<?> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();

core/src/main/java/google/registry/persistence/transaction/TransactionManager.java

@@ -29,17 +29,19 @@ import org.joda.time.DateTime;
  */
 public interface TransactionManager {
 
-  /** Returns {@code true} if the caller is in a transaction.
+  /**
+   * Returns {@code true} if the caller is in a transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   boolean inTransaction();
 
-  /** Throws {@link IllegalStateException} if the caller is not in a transaction.
+  /**
+   * Throws {@link IllegalStateException} if the caller is not in a transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void assertInTransaction();
 
@@ -58,10 +60,11 @@ public interface TransactionManager {
    */
   <T> T transactNew(Supplier<T> work);
 
-  /** Pauses the current transaction (if any) and executes the work in a new transaction.
+  /**
+   * Pauses the current transaction (if any) and executes the work in a new transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void transactNew(Runnable work);
 
@@ -73,10 +76,11 @@ public interface TransactionManager {
    */
   <R> R transactNewReadOnly(Supplier<R> work);
 
-  /** Executes the work in a read-only transaction.
+  /**
+   * Executes the work in a read-only transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void transactNewReadOnly(Runnable work);
 
@@ -182,31 +186,60 @@ public interface TransactionManager {
   /** Returns whether the entity of given key exists. */
   <T> boolean exists(VKey<T> key);
 
-  /** Loads the entity by its id, returns empty if the entity doesn't exist. */
-  <T> Optional<T> maybeLoad(VKey<T> key);
-
-  /** Loads the entity by its id, throws NoSuchElementException if it doesn't exist. */
-  <T> T load(VKey<T> key);
+  /** Loads the entity by its key, returns empty if the entity doesn't exist. */
+  <T> Optional<T> loadByKeyIfPresent(VKey<T> key);
 
   /**
-   * Loads the given entity from the database, throws NoSuchElementException if it doesn't exist.
-   */
-  <T> T load(T entity);
-
-  /**
-   * Loads the set of entities by their key id.
+   * Loads the set of entities by their keys.
    *
-   * @throws NoSuchElementException if any of the keys are not found.
+   * <p>Nonexistent keys / entities are absent from the resulting map, but no {@link
+   * NoSuchElementException} will be thrown.
    */
-  <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys);
-
-  /** Loads all entities of the given type, returns empty if there is no such entity. */
-  <T> ImmutableList<T> loadAll(Class<T> clazz);
+  <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys);
 
   /**
-   * Loads all given entities from the database, throws NoSuchElementException if it doesn't exist.
+   * Loads all given entities from the database if possible.
+   *
+   * <p>Nonexistent entities are absent from the resulting list, but no {@link
+   * NoSuchElementException} will be thrown.
    */
-  <T> ImmutableList<T> loadAll(Iterable<T> entities);
+  <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities);
+
+  /**
+   * Loads the entity by its key.
+   *
+   * @throws NoSuchElementException if this key does not correspond to an existing entity.
+   */
+  <T> T loadByKey(VKey<T> key);
+
+  /**
+   * Loads the set of entities by their keys.
+   *
+   * @throws NoSuchElementException if any of the keys do not correspond to an existing entity.
+   */
+  <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(Iterable<? extends VKey<? extends T>> keys);
+
+  /**
+   * Loads the given entity from the database.
+   *
+   * @throws NoSuchElementException if the entity does not exist in the database.
+   */
+  <T> T loadByEntity(T entity);
+
+  /**
+   * Loads all given entities from the database.
+   *
+   * @throws NoSuchElementException if any of the entities do not exist in the database.
+   */
+  <T> ImmutableList<T> loadByEntities(Iterable<T> entities);
+
+  /**
+   * Returns a stream of all entities of the given type that exist in the database.
+   *
+   * <p>The resulting stream is empty if there are no entities of this type.
+   */
+  <T> ImmutableList<T> loadAllOf(Class<T> clazz);
 
   /** Deletes the entity by its id. */
   void delete(VKey<?> key);

core/src/main/java/google/registry/rdap/RdapJsonFormatter.java

@@ -342,7 +342,7 @@ public class RdapJsonFormatter {
     // Kick off the database loads of the nameservers that we will need, so it can load
     // asynchronously while we load and process the contacts.
     ImmutableSet<HostResource> loadedHosts =
-        ImmutableSet.copyOf(tm().load(domainBase.getNameservers()).values());
+        ImmutableSet.copyOf(tm().loadByKeys(domainBase.getNameservers()).values());
     // Load the registrant and other contacts and add them to the data.
     Map<Key<ContactResource>, ContactResource> loadedContacts =
         ofy()
@@ -429,7 +429,7 @@ public class RdapJsonFormatter {
         statuses.add(StatusValue.LINKED);
       }
       if (hostResource.isSubordinate()
-          && tm().load(hostResource.getSuperordinateDomain())
+          && tm().loadByKey(hostResource.getSuperordinateDomain())
               .cloneProjectedAtTime(getRequestTime())
               .getStatusValues()
               .contains(StatusValue.PENDING_TRANSFER)) {

core/src/main/java/google/registry/rde/DomainBaseToXjcConverter.java

@@ -172,7 +172,7 @@ final class DomainBaseToXjcConverter {
         if (registrant == null) {
           logger.atWarning().log("Domain %s has no registrant contact.", domainName);
         } else {
-          ContactResource registrantContact = tm().load(registrant);
+          ContactResource registrantContact = tm().loadByKey(registrant);
           checkState(
               registrantContact != null,
               "Registrant contact %s on domain %s does not exist",
@@ -305,7 +305,7 @@ final class DomainBaseToXjcConverter {
         "Contact key for type %s is null on domain %s",
         model.getType(),
         domainName);
-    ContactResource contact = tm().load(model.getContactKey());
+    ContactResource contact = tm().loadByKey(model.getContactKey());
     checkState(
         contact != null,
         "Contact %s on domain %s does not exist",

core/src/main/java/google/registry/rde/RdeStagingMapper.java

@@ -203,7 +203,7 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
                         host,
                         // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
                         // us.
-                        loadAtPointInTime(tm().load(host.getSuperordinateDomain()), watermark)
+                        loadAtPointInTime(tm().loadByKey(host.getSuperordinateDomain()), watermark)
                             .now())
                     : marshaller.marshalExternalHost(host));
         cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);

core/src/main/java/google/registry/tools/DeleteAllocationTokensCommand.java

@@ -75,7 +75,7 @@ final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     // Load the tokens in the same transaction as they are deleted to verify they weren't redeemed
     // since the query ran. This also filters out per-domain tokens if they're not to be deleted.
     ImmutableSet<VKey<AllocationToken>> tokensToDelete =
-        tm().load(batch).values().stream()
+        tm().loadByKeys(batch).values().stream()
             .filter(t -> withDomains || t.getDomainName().isEmpty())
             .filter(t -> SINGLE_USE.equals(t.getTokenType()))
             .filter(t -> !t.isRedeemed())

core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -263,7 +263,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
       savedTokens = tokens;
     } else {
       transactIfJpaTm(() -> tm().transact(() -> tm().putAll(tokens)));
-      savedTokens = tm().transact(() -> tm().loadAll(tokens));
+      savedTokens = tm().transact(() -> tm().loadByEntities(tokens));
     }
     savedTokens.forEach(
         t -> System.out.println(SKIP_NULLS.join(t.getDomainName().orElse(null), t.getToken())));
@@ -293,7 +293,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
             .collect(toImmutableSet());
     return transactIfJpaTm(
         () ->
-            tm().load(existingTokenKeys).values().stream()
+            tm().loadByKeysIfPresent(existingTokenKeys).values().stream()
                 .map(AllocationToken::getToken)
                 .collect(toImmutableSet()));
   }

core/src/main/java/google/registry/tools/GetAllocationTokenCommand.java

@@ -49,7 +49,8 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
           tokens.stream()
               .map(t -> VKey.create(AllocationToken.class, t))
               .collect(toImmutableList());
-      tm().load(tokenKeys).forEach((k, v) -> builder.put(k.getSqlKey().toString(), v));
+      tm().loadByKeysIfPresent(tokenKeys)
+          .forEach((k, v) -> builder.put(k.getSqlKey().toString(), v));
     }
     ImmutableMap<String, AllocationToken> loadedTokens = builder.build();
     ImmutableMap<VKey<DomainBase>, DomainBase> domains = loadRedeemedDomains(loadedTokens.values());
@@ -88,14 +89,14 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
             .map(AllocationToken::getRedemptionHistoryEntry)
             .filter(Optional::isPresent)
             .map(Optional::get)
-            .map(key -> tm().load(key))
+            .map(key -> tm().loadByKey(key))
             .map(he -> (Key<DomainBase>) he.getParent())
             .map(key -> VKey.create(DomainBase.class, key.getName(), key))
             .collect(toImmutableList());
     ImmutableMap.Builder<VKey<DomainBase>, DomainBase> domainsBuilder =
         new ImmutableMap.Builder<>();
     for (List<VKey<DomainBase>> keys : Lists.partition(domainKeys, BATCH_SIZE)) {
-      tm().load(ImmutableList.copyOf(keys))
+      tm().loadByKeys(ImmutableList.copyOf(keys))
           .forEach((k, v) -> domainsBuilder.put((VKey<DomainBase>) k, v));
     }
     return domainsBuilder.build();

core/src/main/java/google/registry/tools/ResaveEppResourceCommand.java

@@ -56,7 +56,7 @@ public final class ResaveEppResourceCommand extends MutatingCommand {
             uniqueId);
     // Load the resource directly to bypass running cloneProjectedAtTime() automatically, which can
     // cause stageEntityChange() to fail due to implicit projection changes.
-    EppResource resource = tm().load(resourceKey);
+    EppResource resource = tm().loadByKey(resourceKey);
     stageEntityChange(resource, resource);
   }
 }

core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java

@@ -162,7 +162,7 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
 
   private ImmutableSortedSet<String> getExistingNameservers(DomainBase domain) {
     ImmutableSortedSet.Builder<String> nameservers = ImmutableSortedSet.naturalOrder();
-    for (HostResource host : tm().load(domain.getNameservers()).values()) {
+    for (HostResource host : tm().loadByKeys(domain.getNameservers()).values()) {
       nameservers.add(host.getForeignKey());
     }
     return nameservers.build();

core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java

@@ -107,7 +107,7 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     tokensToSave =
         transactIfJpaTm(
             () ->
-                tm().load(getTokenKeys()).values().stream()
+                tm().loadByKeys(getTokenKeys()).values().stream()
                     .collect(toImmutableMap(Function.identity(), this::updateToken))
                     .entrySet()
                     .stream()

core/src/main/java/google/registry/tools/UpdateDomainCommand.java

@@ -329,7 +329,7 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
       DomainBase domainBase, final DesignatedContact.Type contactType) {
     return domainBase.getContacts().stream()
         .filter(contact -> contact.getType().equals(contactType))
-        .map(contact -> tm().load(contact.getContactKey()).getContactId())
+        .map(contact -> tm().loadByKey(contact.getContactKey()).getContactId())
         .collect(toImmutableSet());
   }
 }

core/src/main/java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java

@@ -67,7 +67,7 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand
       checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
       return transactIfJpaTm(
           () ->
-              tm().loadAll(AllocationToken.class).stream()
+              tm().loadAllOf(AllocationToken.class).stream()
                   .filter(token -> token.getToken().startsWith(prefix))
                   .map(AllocationToken::createVKey)
                   .collect(toImmutableSet()));

core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java

@@ -214,7 +214,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
     private void emitForSubordinateHosts(DomainBase domain) {
       ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
       if (!subordinateHosts.isEmpty()) {
-        for (HostResource unprojectedHost : tm().load(domain.getNameservers()).values()) {
+        for (HostResource unprojectedHost : tm().loadByKeys(domain.getNameservers()).values()) {
           HostResource host = loadAtPointInTime(unprojectedHost, exportTime).now();
           // A null means the host was deleted (or not created) at this time.
           if ((host != null) && subordinateHosts.contains(host.getHostName())) {
@@ -283,7 +283,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
       Duration dnsDefaultDsTtl) {
     StringBuilder result = new StringBuilder();
     String domainLabel = stripTld(domain.getDomainName(), domain.getTld());
-    for (HostResource nameserver : tm().load(domain.getNameservers()).values()) {
+    for (HostResource nameserver : tm().loadByKeys(domain.getNameservers()).values()) {
       result.append(
           String.format(
               NS_FORMAT,

core/src/main/java/google/registry/whois/NameserverWhoisResponse.java

@@ -49,7 +49,7 @@ final class NameserverWhoisResponse extends WhoisResponseImpl {
       HostResource host = hosts.get(i);
       String clientId =
           host.isSubordinate()
-              ? tm().load(host.getSuperordinateDomain())
+              ? tm().loadByKey(host.getSuperordinateDomain())
                   .cloneProjectedAtTime(getTimestamp())
                   .getCurrentSponsorClientId()
               : host.getPersistedCurrentSponsorClientId();