From 10e82953aeafc3e2fd552452cbbce276fa5b7c0b Mon Sep 17 00:00:00 2001 From: Michael Muller Date: Fri, 26 Jun 2020 13:50:02 -0400 Subject: [PATCH] Make EppResource.loadCached() use batched fetch (#652) * Make EppResource.loadCached() use batched fetch Use a batched fetch (ofy().load().keys(...)) from datastore in EppResource.loadCached(). To support this, convert TransactionManager.load(Iterable) to accept the more flexible generic parameters and return a map. * Simplify datastore key streaming * Changes requested in review. --- .../google/registry/model/EppResource.java | 19 ++-------- .../ofy/DatastoreTransactionManager.java | 18 ++++++---- .../JpaTransactionManagerImpl.java | 23 ++++++------ .../transaction/TransactionManager.java | 3 +- .../registry/rdap/RdapJsonFormatter.java | 2 +- .../tools/UniformRapidSuspensionCommand.java | 2 +- .../tools/server/GenerateZoneFilesAction.java | 4 +-- .../transaction/TransactionManagerTest.java | 36 +++++++++++++++++++ 8 files changed, 68 insertions(+), 39 deletions(-) diff --git a/core/src/main/java/google/registry/model/EppResource.java b/core/src/main/java/google/registry/model/EppResource.java index e2ccc25af..2c10ad714 100644 --- a/core/src/main/java/google/registry/model/EppResource.java +++ b/core/src/main/java/google/registry/model/EppResource.java @@ -16,7 +16,6 @@ package google.registry.model; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; -import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.Sets.difference; import static com.google.common.collect.Sets.union; import static google.registry.config.RegistryConfig.getEppResourceCachingDuration; @@ -47,7 +46,6 @@ import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.ExecutionException; -import java.util.stream.StreamSupport; import javax.persistence.Access; import javax.persistence.AccessType; import javax.persistence.Column; @@ -357,7 +355,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable { @Override public Map, EppResource> loadAll( Iterable> keys) { - return tm().doTransactionless(() -> loadAsMap(keys)); + return tm().doTransactionless(() -> tm().load(keys)); } }; @@ -397,7 +395,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable { public static ImmutableMap, EppResource> loadCached( Iterable> keys) { if (!RegistryConfig.isEppResourceCachingEnabled()) { - return loadAsMap(keys); + return tm().load(keys); } try { return cacheEppResources.getAll(keys); @@ -425,17 +423,4 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable { throw new RuntimeException("Error loading cached EppResources", e.getCause()); } } - - private static ImmutableMap, EppResource> loadAsMap( - Iterable> keys) { - return StreamSupport.stream(keys.spliterator(), false) - // It's possible for us to receive the same key more than once which causes - // the immutable map build to break with a duplicate key, so we have to ensure key - // uniqueness. - .distinct() - // We have to use "key -> key" here instead of the identity() function, because - // the latter breaks the fairly complicated generic type checking required by the - // caching interface. - .collect(toImmutableMap(key -> key, key -> tm().load(key))); - } } diff --git a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java index 07b59328d..723cc0703 100644 --- a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java +++ b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java @@ -15,15 +15,18 @@ package google.registry.model.ofy; import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static google.registry.model.ofy.ObjectifyService.ofy; import static google.registry.util.PreconditionsUtils.checkArgumentNotNull; +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.googlecode.objectify.Key; import google.registry.persistence.VKey; import google.registry.persistence.transaction.TransactionManager; -import java.util.Iterator; +import java.util.Map.Entry; import java.util.NoSuchElementException; import java.util.Optional; import java.util.function.Supplier; @@ -156,12 +159,15 @@ public class DatastoreTransactionManager implements TransactionManager { } @Override - public ImmutableList load(Iterable> keys) { - Iterator> iter = - StreamSupport.stream(keys.spliterator(), false).map(VKey::getOfyKey).iterator(); + public ImmutableMap, T> load(Iterable> keys) { + // Keep track of the Key -> VKey mapping so we can translate them back. + ImmutableMap, VKey> keyMap = + StreamSupport.stream(keys.spliterator(), false) + .distinct() + .collect(toImmutableMap(key -> (Key) key.getOfyKey(), Functions.identity())); - // The lambda argument to keys() effectively converts Iterator -> Iterable. - return ImmutableList.copyOf(getOfy().load().keys(() -> iter).values()); + return getOfy().load().keys(keyMap.keySet()).entrySet().stream() + .collect(ImmutableMap.toImmutableMap(entry -> keyMap.get(entry.getKey()), Entry::getValue)); } @Override diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index 22e6d55ca..b1f2cf062 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java @@ -15,18 +15,21 @@ package google.registry.persistence.transaction; import static com.google.common.base.Preconditions.checkArgument; -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.util.PreconditionsUtils.checkArgumentNotNull; +import static java.util.AbstractMap.SimpleEntry; import static java.util.stream.Collectors.joining; 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.flogger.FluentLogger; import google.registry.persistence.VKey; import google.registry.util.Clock; import java.lang.reflect.Field; +import java.util.Map.Entry; import java.util.NoSuchElementException; import java.util.Optional; import java.util.function.Supplier; @@ -253,20 +256,18 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { } @Override - public ImmutableList load(Iterable> keys) { + public ImmutableMap, T> load(Iterable> keys) { checkArgumentNotNull(keys, "keys must be specified"); assertInTransaction(); return StreamSupport.stream(keys.spliterator(), false) + // Accept duplicate keys. + .distinct() .map( - key -> { - T entity = getEntityManager().find(key.getKind(), key.getSqlKey()); - if (entity == null) { - throw new NoSuchElementException( - key.getKind().getName() + " with key " + key.getSqlKey() + " not found."); - } - return entity; - }) - .collect(toImmutableList()); + key -> + new SimpleEntry, T>( + key, getEntityManager().find(key.getKind(), key.getSqlKey()))) + .filter(entry -> entry.getValue() != null) + .collect(toImmutableMap(Entry::getKey, Entry::getValue)); } @Override diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java index 3c522891e..89f453959 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java +++ b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java @@ -16,6 +16,7 @@ package google.registry.persistence.transaction; import com.google.common.collect.ImmutableCollection; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import google.registry.persistence.VKey; import java.util.NoSuchElementException; import java.util.Optional; @@ -119,7 +120,7 @@ public interface TransactionManager { * * @throws NoSuchElementException if any of the keys are not found. */ - ImmutableList load(Iterable> keys); + ImmutableMap, T> load(Iterable> keys); /** Loads all entities of the given type, returns empty if there is no such entity. */ ImmutableList loadAll(Class clazz); diff --git a/core/src/main/java/google/registry/rdap/RdapJsonFormatter.java b/core/src/main/java/google/registry/rdap/RdapJsonFormatter.java index 11d2cf2d0..9fbb33758 100644 --- a/core/src/main/java/google/registry/rdap/RdapJsonFormatter.java +++ b/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 loadedHosts = - ImmutableSet.copyOf(tm().load(domainBase.getNameservers())); + ImmutableSet.copyOf(tm().load(domainBase.getNameservers()).values()); // Load the registrant and other contacts and add them to the data. Map, ContactResource> loadedContacts = ofy() diff --git a/core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java b/core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java index 025e044c1..2eb2c51cc 100644 --- a/core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java +++ b/core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java @@ -149,7 +149,7 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand { private ImmutableSortedSet getExistingNameservers(DomainBase domain) { ImmutableSortedSet.Builder nameservers = ImmutableSortedSet.naturalOrder(); - for (HostResource host : tm().load(domain.getNameservers())) { + for (HostResource host : tm().load(domain.getNameservers()).values()) { nameservers.add(host.getForeignKey()); } return nameservers.build(); diff --git a/core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java b/core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java index e0410adea..aa64ccc1e 100644 --- a/core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java +++ b/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 subordinateHosts = domain.getSubordinateHosts(); if (!subordinateHosts.isEmpty()) { - for (HostResource unprojectedHost : tm().load(domain.getNameservers())) { + for (HostResource unprojectedHost : tm().load(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())) { + for (HostResource nameserver : tm().load(domain.getNameservers()).values()) { result.append( String.format( NS_FORMAT, diff --git a/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java b/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java index 76adb56c4..262859d97 100644 --- a/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java +++ b/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java @@ -21,6 +21,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory. import static org.junit.Assert.assertThrows; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Maps; import com.googlecode.objectify.Key; import com.googlecode.objectify.annotation.Entity; import com.googlecode.objectify.annotation.Id; @@ -35,6 +36,7 @@ import google.registry.testing.InjectRule; import java.util.List; import java.util.NoSuchElementException; import java.util.Set; +import java.util.stream.Stream; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.TestTemplate; import org.junit.jupiter.api.extension.RegisterExtension; @@ -272,6 +274,40 @@ public class TransactionManagerTest { assertAllEntitiesNotExist(moreEntities); } + @TestTemplate + void load_multi() { + assertAllEntitiesNotExist(moreEntities); + tm().transact(() -> tm().saveAllNew(moreEntities)); + List> keys = + moreEntities.stream().map(TestEntity::key).collect(toImmutableList()); + assertThat(tm().transact(() -> tm().load(keys))) + .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key)); + } + + @TestTemplate + void load_multiWithDuplicateKeys() { + assertAllEntitiesNotExist(moreEntities); + tm().transact(() -> tm().saveAllNew(moreEntities)); + ImmutableList> keys = + moreEntities.stream().map(TestEntity::key).collect(toImmutableList()); + ImmutableList> doubleKeys = + Stream.concat(keys.stream(), keys.stream()).collect(toImmutableList()); + assertThat(tm().transact(() -> tm().load(doubleKeys))) + .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key)); + } + + @TestTemplate + void load_multiMissingKeys() { + assertAllEntitiesNotExist(moreEntities); + tm().transact(() -> tm().saveAllNew(moreEntities)); + List> keys = + Stream.concat(moreEntities.stream(), Stream.of(new TestEntity("dark", "matter"))) + .map(TestEntity::key) + .collect(toImmutableList()); + assertThat(tm().transact(() -> tm().load(keys))) + .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key)); + } + private static void assertEntityExists(TestEntity entity) { assertThat(tm().transact(() -> tm().checkExists(entity))).isTrue(); }