diff --git a/core/src/main/java/google/registry/beam/rde/RdePipeline.java b/core/src/main/java/google/registry/beam/rde/RdePipeline.java index 37507988e1b..88d34ec8802 100644 --- a/core/src/main/java/google/registry/beam/rde/RdePipeline.java +++ b/core/src/main/java/google/registry/beam/rde/RdePipeline.java @@ -16,6 +16,8 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +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.beam.rde.RdePipeline.TupleTags.DOMAIN_FRAGMENTS; import static google.registry.beam.rde.RdePipeline.TupleTags.EXTERNAL_HOST_FRAGMENTS; @@ -24,12 +26,12 @@ import static google.registry.beam.rde.RdePipeline.TupleTags.REFERENCED_HOSTS; import static google.registry.beam.rde.RdePipeline.TupleTags.REVISION_ID; import static google.registry.beam.rde.RdePipeline.TupleTags.SUPERORDINATE_DOMAINS; -import static google.registry.model.reporting.HistoryEntryDao.RESOURCE_TYPES_TO_HISTORY_TYPES; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.util.SafeSerializationUtils.safeDeserializeCollection; import static google.registry.util.SafeSerializationUtils.serializeCollection; import static google.registry.util.SerializeUtils.decodeBase64; import static google.registry.util.SerializeUtils.encodeBase64; +import static org.apache.beam.sdk.values.TypeDescriptors.integers; import static org.apache.beam.sdk.values.TypeDescriptors.kvs; import com.google.common.collect.ImmutableList; @@ -57,7 +59,6 @@ import google.registry.model.registrar.Registrar; import google.registry.model.registrar.Registrar.Type; import google.registry.model.reporting.HistoryEntry; -import google.registry.model.reporting.HistoryEntry.HistoryEntryId; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; import google.registry.persistence.VKey; import google.registry.rde.DepositFragment; @@ -71,6 +72,7 @@ import java.io.IOException; import java.io.Serializable; import java.time.Instant; +import java.util.NoSuchElementException; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.KvCoder; @@ -85,10 +87,13 @@ import org.apache.beam.sdk.transforms.FlatMapElements; import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.GroupIntoBatches; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.join.CoGbkResult; import org.apache.beam.sdk.transforms.join.CoGroupByKey; import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; @@ -327,8 +332,8 @@ private PCollection> getMostRecentHist .withCoder(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))); } - private EppResource loadResourceByHistoryEntryId( - Class historyEntryClazz, String repoId, Iterable revisionIds) { + private static long getSingleRevisionId( + Class historyEntryClazz, String repoId, Iterable revisionIds) { ImmutableList ids = ImmutableList.copyOf(revisionIds); // The size should always be 1 because we are only getting one repo ID -> revision ID pair per // repo ID from the source transform (the JPA query in the method above). But for some reason @@ -351,22 +356,52 @@ private EppResource loadResourceByHistoryEntryId( "Duplicate revision IDs detected for %s repo ID %s: %s", historyEntryClazz.getSimpleName(), repoId, ids); } - return loadResourceByHistoryEntryId(historyEntryClazz, repoId, ids.get(0)); + return ids.getFirst(); } - private EppResource loadResourceByHistoryEntryId( - Class historyEntryClazz, String repoId, long revisionId) { - return tm().transact( - () -> - tm().loadByKey( - VKey.create(historyEntryClazz, new HistoryEntryId(repoId, revisionId)))) - .getResourceAtPointInTime() - .map(resource -> resource.cloneProjectedAtTime(watermark)) - .get(); + + ImmutableMap loadResourcesByHistoryEntryIds( + Iterable> repoAndRevisionIds, + Class resourceClass, + Class historyEntryClass) { + ImmutableList> ids = ImmutableList.copyOf(repoAndRevisionIds); + if (ids.isEmpty()) { + return ImmutableMap.of(); + } + ImmutableList revisionIds = ids.stream().map(KV::getValue).collect(toImmutableList()); + ImmutableMap result = + tm().transact( + () -> + tm().query( + String.format( + "FROM %s WHERE revisionId IN (:revisionIds)", + historyEntryClass.getSimpleName()), + historyEntryClass) + .setParameter("revisionIds", revisionIds) + .getResultStream() + .collect( + toImmutableMap( + HistoryEntry::getRepoId, + entry -> + entry + .getResourceAtPointInTime() + .map(r -> r.cloneProjectedAtTime(watermark)) + .map(resourceClass::cast) + .get()))); + // Fail fast on items being missing unexpectedly + if (result.size() != ids.size()) { + throw new NoSuchElementException( + String.format( + "Expected to find the following %s history entries but they were missing: %s", + historyEntryClass.getSimpleName(), + Sets.difference( + ids.stream().map(KV::getKey).collect(toImmutableSet()), result.keySet()))); + } + return result; } /** - * Remove unreferenced resources by joining the (repoId, pendingDeposit) pair with the (repoId, + * Remove unreferenced hosts by joining the (repoId, pendingDeposit) pair with the (repoId, * revisionId) on the repoId. * *

The (repoId, pendingDeposit) pairs denote hosts that are referenced from a domain, that are @@ -380,21 +415,13 @@ private EppResource loadResourceByHistoryEntryId( */ private static PCollection> removeUnreferencedResource( PCollection> referencedResources, - PCollection> historyEntries, - Class resourceClazz) { - String resourceName = resourceClazz.getSimpleName(); - Class historyEntryClazz = - RESOURCE_TYPES_TO_HISTORY_TYPES.get(resourceClazz); - String historyEntryName = historyEntryClazz.getSimpleName(); - Counter referencedResourceCounter = Metrics.counter("RDE", "Referenced" + resourceName); + PCollection> hostHistories) { + Counter referencedResourceCounter = Metrics.counter("RDE", "Referenced Host"); return KeyedPCollectionTuple.of(PENDING_DEPOSIT, referencedResources) - .and(REVISION_ID, historyEntries) + .and(REVISION_ID, hostHistories) + .apply("Join PendingDeposit with HostHistory revision ID on Host", CoGroupByKey.create()) .apply( - String.format( - "Join PendingDeposit with %s revision ID on %s", historyEntryName, resourceName), - CoGroupByKey.create()) - .apply( - String.format("Remove unreferenced %s", resourceName), + "Remove unreferenced Host", Filter.by( (KV kv) -> { boolean toInclude = @@ -419,50 +446,65 @@ private PCollectionTuple processDomainHistories(PCollection> do Counter activeDomainCounter = Metrics.counter("RDE", "ActiveDomainBase"); Counter domainFragmentCounter = Metrics.counter("RDE", "DomainFragment"); Counter referencedHostCounter = Metrics.counter("RDE", "ReferencedHost"); - return domainHistories.apply( - "Map DomainHistory to DepositFragment and emit referenced Host", - ParDo.of( - new DoFn, KV>() { - @ProcessElement - public void processElement( - @Element KV kv, MultiOutputReceiver receiver) { - activeDomainCounter.inc(); - Domain domain = - (Domain) - loadResourceByHistoryEntryId( - DomainHistory.class, kv.getKey(), kv.getValue()); - pendingDeposits.stream() - .filter(pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) - .forEach( - pendingDeposit -> { - // Domains are always deposited in both modes. - domainFragmentCounter.inc(); - receiver - .get(DOMAIN_FRAGMENTS) - .output( - KV.of( - pendingDeposit, - marshaller.marshalDomain(domain, pendingDeposit.mode()))); - // Hosts are only deposited in RDE, not BRDA. - if (pendingDeposit.mode() == RdeMode.FULL) { - if (domain.getNsHosts() != null) { - referencedHostCounter.inc(domain.getNsHosts().size()); - domain - .getNsHosts() - .forEach( - hostKey -> - receiver - .get(REFERENCED_HOSTS) - .output( - KV.of( - (String) hostKey.getKey(), - pendingDeposit))); - } - } - }); - } - }) - .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); + return domainHistories + .apply( + "Assign all domain histories a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group domain histories into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) + .apply( + "Map DomainHistory to DepositFragment and emit referenced Host", + ParDo.of( + new DoFn< + KV, Iterable>>, + KV>() { + @ProcessElement + public void processElement( + @Element KV, Iterable>> element, + MultiOutputReceiver receiver) { + ImmutableMap loadedDomains = + loadResourcesByHistoryEntryIds( + element.getValue(), Domain.class, DomainHistory.class); + for (KV kv : element.getValue()) { + Domain domain = loadedDomains.get(kv.getKey()); + activeDomainCounter.inc(); + pendingDeposits.stream() + .filter( + pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) + .forEach( + pendingDeposit -> { + // Domains are always deposited in both modes. + domainFragmentCounter.inc(); + receiver + .get(DOMAIN_FRAGMENTS) + .output( + KV.of( + pendingDeposit, + marshaller.marshalDomain( + domain, pendingDeposit.mode()))); + // Hosts are only deposited in RDE, not BRDA. + if (pendingDeposit.mode() == RdeMode.FULL) { + if (domain.getNsHosts() != null) { + referencedHostCounter.inc(domain.getNsHosts().size()); + domain + .getNsHosts() + .forEach( + hostKey -> + receiver + .get(REFERENCED_HOSTS) + .output( + KV.of( + (String) hostKey.getKey(), + pendingDeposit))); + } + } + }); + } + } + }) + .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); } private PCollectionTuple processHostHistories( @@ -471,45 +513,66 @@ private PCollectionTuple processHostHistories( Counter subordinateHostCounter = Metrics.counter("RDE", "SubordinateHost"); Counter externalHostCounter = Metrics.counter("RDE", "ExternalHost"); Counter externalHostFragmentCounter = Metrics.counter("RDE", "ExternalHostFragment"); - return removeUnreferencedResource(referencedHosts, hostHistories, Host.class) + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); + return removeUnreferencedResource(referencedHosts, hostHistories) + .apply( + "Assign all host histories a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group referenced hosts into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) .apply( "Map external DomainResource to DepositFragment and process subordinate domains", ParDo.of( - new DoFn, KV>() { + new DoFn< + KV, Iterable>>, + KV>() { @ProcessElement public void processElement( - @Element KV kv, MultiOutputReceiver receiver) { - Host host = - (Host) - loadResourceByHistoryEntryId( - HostHistory.class, - kv.getKey(), - kv.getValue().getAll(REVISION_ID)); - // When a host is subordinate, we need to find its superordinate domain and - // include it in the deposit as well. - if (host.isSubordinate()) { - subordinateHostCounter.inc(); - receiver - .get(SUPERORDINATE_DOMAINS) - .output( - // The output are pairs of - // (superordinateDomainRepoId, - // (subordinateHostRepoId, (pendingDeposit, revisionId))). - KV.of((String) host.getSuperordinateDomain().getKey(), kv)); - } else { - externalHostCounter.inc(); - DepositFragment fragment = marshaller.marshalExternalHost(host); - Streams.stream(kv.getValue().getAll(PENDING_DEPOSIT)) - // The same host could be used by multiple domains, therefore - // matched to the same pending deposit multiple times. - .distinct() - .forEach( - pendingDeposit -> { - externalHostFragmentCounter.inc(); - receiver - .get(EXTERNAL_HOST_FRAGMENTS) - .output(KV.of(pendingDeposit, fragment)); - }); + @Element + KV, Iterable>> element, + MultiOutputReceiver receiver) { + ImmutableSet> hostKeys = + Streams.stream(element.getValue()) + .map( + kv -> + KV.of( + kv.getKey(), + getSingleRevisionId( + HostHistory.class, + kv.getKey(), + kv.getValue().getAll(REVISION_ID)))) + .collect(toImmutableSet()); + ImmutableMap loadedHosts = + loadResourcesByHistoryEntryIds(hostKeys, Host.class, HostHistory.class); + for (KV kv : element.getValue()) { + Host host = loadedHosts.get(kv.getKey()); + // When a host is subordinate, we need to find its superordinate domain + // and include it in the deposit as well. + if (host.isSubordinate()) { + subordinateHostCounter.inc(); + receiver + .get(SUPERORDINATE_DOMAINS) + .output( + // The output are pairs of + // (superordinateDomainRepoId, + // (subordinateHostRepoId, (pendingDeposit, revisionId))). + KV.of((String) host.getSuperordinateDomain().getKey(), kv)); + } else { + externalHostCounter.inc(); + DepositFragment fragment = marshaller.marshalExternalHost(host); + Streams.stream(kv.getValue().getAll(PENDING_DEPOSIT)) + // The same host could be used by multiple domains, therefore + // matched to the same pending deposit multiple times. + .distinct() + .forEach( + pendingDeposit -> { + externalHostFragmentCounter.inc(); + receiver + .get(EXTERNAL_HOST_FRAGMENTS) + .output(KV.of(pendingDeposit, fragment)); + }); + } } } }) @@ -532,6 +595,7 @@ private PCollection> processSubordinateHosts PCollection> domainHistories) { Counter subordinateHostFragmentCounter = Metrics.counter("RDE", "SubordinateHostFragment"); Counter referencedSubordinateHostCounter = Metrics.counter("RDE", "ReferencedSubordinateHost"); + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT, superordinateDomains) .and(REVISION_ID, domainHistories) .apply("Join Host:PendingDeposits with DomainHistory on Domain", CoGroupByKey.create()) @@ -547,30 +611,59 @@ private PCollection> processSubordinateHosts } return toInclude; })) + .apply( + "Assign all subordinate hosts a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group subordinate hosts into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) .apply( "Map subordinate Host to DepositFragment", - FlatMapElements.into( - kvs( - TypeDescriptor.of(PendingDeposit.class), - TypeDescriptor.of(DepositFragment.class))) - .via( - (KV kv) -> { - Domain superordinateDomain = - (Domain) - loadResourceByHistoryEntryId( - DomainHistory.class, - kv.getKey(), - kv.getValue().getAll(REVISION_ID)); - ImmutableSet.Builder> results = - new ImmutableSet.Builder<>(); + ParDo.of( + new DoFn< + KV, Iterable>>, + KV>() { + @ProcessElement + public void processElement( + @Element KV, Iterable>> element, + OutputReceiver> receiver) { + ImmutableSet> domainKeys = + Streams.stream(element.getValue()) + .map( + kv -> + KV.of( + kv.getKey(), + getSingleRevisionId( + DomainHistory.class, + kv.getKey(), + kv.getValue().getAll(REVISION_ID)))) + .collect(toImmutableSet()); + ImmutableSet> hostKeys = + Streams.stream(element.getValue()) + .flatMap( + kv -> + Streams.stream(kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT)) + .map( + hostToPendingDeposits -> + KV.of( + hostToPendingDeposits.getKey(), + getSingleRevisionId( + HostHistory.class, + hostToPendingDeposits.getKey(), + hostToPendingDeposits + .getValue() + .getAll(REVISION_ID))))) + .collect(toImmutableSet()); + ImmutableMap loadedDomains = + loadResourcesByHistoryEntryIds( + domainKeys, Domain.class, DomainHistory.class); + ImmutableMap loadedHosts = + loadResourcesByHistoryEntryIds(hostKeys, Host.class, HostHistory.class); + for (KV kv : element.getValue()) { + Domain superordinateDomain = loadedDomains.get(kv.getKey()); for (KV hostToPendingDeposits : kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT)) { - Host host = - (Host) - loadResourceByHistoryEntryId( - HostHistory.class, - hostToPendingDeposits.getKey(), - hostToPendingDeposits.getValue().getAll(REVISION_ID)); + Host host = loadedHosts.get(hostToPendingDeposits.getKey()); DepositFragment fragment = marshaller.marshalSubordinateHost(host, superordinateDomain); Streams.stream(hostToPendingDeposits.getValue().getAll(PENDING_DEPOSIT)) @@ -578,11 +671,12 @@ private PCollection> processSubordinateHosts .forEach( pendingDeposit -> { subordinateHostFragmentCounter.inc(); - results.add(KV.of(pendingDeposit, fragment)); + receiver.output(KV.of(pendingDeposit, fragment)); }); } - return results.build(); - })); + } + } + })); } /** diff --git a/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java b/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java index c365494b67b..e3601e198d2 100644 --- a/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java +++ b/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java @@ -15,9 +15,10 @@ package google.registry.beam.rde; import google.registry.beam.common.RegistryPipelineOptions; +import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; -/** Custom options for running the spec11 pipeline. */ +/** Custom options for running the RDE pipeline. */ public interface RdePipelineOptions extends RegistryPipelineOptions { @Description("The Base64-encoded serialized map of TLDs to PendingDeposit.") @@ -39,4 +40,11 @@ public interface RdePipelineOptions extends RegistryPipelineOptions { String getStagingKey(); void setStagingKey(String value); + + @Description( + "The number of history entries to batch load from the SQL database in one operation.") + @Default.Integer(500) + int getRdeHistoryEntryLoadBatchSize(); + + void setRdeHistoryEntryLoadBatchSize(int value); } diff --git a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java index bd8fa1314fc..de2ff0c6a5a 100644 --- a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java +++ b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java @@ -85,6 +85,7 @@ import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.NoSuchElementException; import java.util.function.Function; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -166,36 +167,42 @@ private DomainHistory persistDomainHistory(DomainBase domain) { .setReportAmount(1) .build(); - return persistResource( - new DomainHistory.Builder() - .setType(HistoryEntry.Type.DOMAIN_CREATE) - .setXmlBytes("".getBytes(UTF_8)) - .setModificationTime(clock.now()) - .setRegistrarId("TheRegistrar") - .setTrid(Trid.create("ABC-123", "server-trid")) - .setBySuperuser(false) - .setReason("reason") - .setRequestedByRegistrar(true) - .setDomain(domain) - .setDomainTransactionRecords(ImmutableSet.of(transactionRecord)) - .setOtherRegistrarId("otherClient") - .setPeriod(Period.create(1, Period.Unit.YEARS)) - .build()); + DomainHistory result = + persistResource( + new DomainHistory.Builder() + .setType(HistoryEntry.Type.DOMAIN_CREATE) + .setXmlBytes("".getBytes(UTF_8)) + .setModificationTime(clock.now()) + .setRegistrarId("TheRegistrar") + .setTrid(Trid.create("ABC-123", "server-trid")) + .setBySuperuser(false) + .setReason("reason") + .setRequestedByRegistrar(true) + .setDomain(domain) + .setDomainTransactionRecords(ImmutableSet.of(transactionRecord)) + .setOtherRegistrarId("otherClient") + .setPeriod(Period.create(1, Period.Unit.YEARS)) + .build()); + clock.advanceOneMilli(); + return result; } private HostHistory persistHostHistory(HostBase hostBase) { - return persistResource( - new HostHistory.Builder() - .setType(HistoryEntry.Type.HOST_CREATE) - .setXmlBytes("".getBytes(UTF_8)) - .setModificationTime(clock.now()) - .setRegistrarId("TheRegistrar") - .setTrid(Trid.create("ABC-123", "server-trid")) - .setBySuperuser(false) - .setReason("reason") - .setRequestedByRegistrar(true) - .setHost(hostBase) - .build()); + HostHistory result = + persistResource( + new HostHistory.Builder() + .setType(HistoryEntry.Type.HOST_CREATE) + .setXmlBytes("".getBytes(UTF_8)) + .setModificationTime(clock.now()) + .setRegistrarId("TheRegistrar") + .setTrid(Trid.create("ABC-123", "server-trid")) + .setBySuperuser(false) + .setReason("reason") + .setRequestedByRegistrar(true) + .setHost(hostBase) + .build()); + clock.advanceOneMilli(); + return result; } @BeforeEach @@ -257,7 +264,7 @@ void beforeEach() throws Exception { Domain deletedDomain = persistActiveDomain("deleted.soy"); persistDomainHistory(deletedDomain); - // Advance time + // Advance time again just in case clock.advanceOneMilli(); persistDomainHistory(deletedDomain.asBuilder().setDeletionTime(clock.now()).build()); kittyDomain = kittyDomain.asBuilder().setDomainName("cat.fun").build(); @@ -425,6 +432,31 @@ void testSuccess_createFragments() { pipeline.run().waitUntilFinish(); } + @Test + void testSuccess_createFragments_smallBatchSize() { + options.setRdeHistoryEntryLoadBatchSize(1); + testSuccess_createFragments(); + } + + @Test + void testSuccess_createFragments_multiBatch() { + options.setRdeHistoryEntryLoadBatchSize(2); + testSuccess_createFragments(); + } + + @Test + void testFailure_missingHistoryEntry() { + NoSuchElementException thrown = + assertThrows( + NoSuchElementException.class, + () -> + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of(KV.of("nonexistent", 12345L)), + Domain.class, + DomainHistory.class)); + assertThat(thrown).hasMessageThat().contains("nonexistent"); + } + // The GCS folder listing can be a bit flaky, so retry if necessary @RetryingTest(4) void testSuccess_persistData() throws Exception {