From a8f01d590d9dceaa4cb57cf952d06c6d0f2ead17 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Sat, 15 Aug 2026 01:20:27 +0000 Subject: [PATCH] Batch database queries in RdePipeline (#3190) Currently, we process (repoId, revisionId) pairs for DomainHistory and HostHistory individually -- they may be farmed out to worker nodes in parallel, but each EppResource uses a separate transaction and a separate read, which doesn't scale well when there are lots of domains/hosts. So as a result, we should batch them up so we can load (by default) 500 per transaction at a time. We don't want to batch-load the domains/hosts at the same time that we retrieve the most recent history entry for each type -- this would mean passing relatively large objects across pipeline steps. Instead, we keep passing the KV and batch retrievals. This isn't necessarily much faster (due to having to wait on batching) but there'll be less load on the DB. Self-scan D.2 number 5 --- .../google/registry/beam/rde/RdePipeline.java | 554 +++++++++++------- .../registry/beam/rde/RdePipelineOptions.java | 16 +- .../registry/beam/rde_pipeline_metadata.json | 18 + .../registry/beam/rde/RdePipelineTest.java | 141 ++++- 4 files changed, 497 insertions(+), 232 deletions(-) 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 37507988e..6dcd3ef23 100644 --- a/core/src/main/java/google/registry/beam/rde/RdePipeline.java +++ b/core/src/main/java/google/registry/beam/rde/RdePipeline.java @@ -16,25 +16,27 @@ package google.registry.beam.rde; 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.ImmutableSet.toImmutableSet; import static google.registry.beam.rde.RdePipeline.TupleTags.DOMAIN_FRAGMENTS; import static google.registry.beam.rde.RdePipeline.TupleTags.EXTERNAL_HOST_FRAGMENTS; -import static google.registry.beam.rde.RdePipeline.TupleTags.HOST_TO_PENDING_DEPOSIT; +import static google.registry.beam.rde.RdePipeline.TupleTags.HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID; import static google.registry.beam.rde.RdePipeline.TupleTags.PENDING_DEPOSIT; 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; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import com.google.common.collect.Streams; import com.google.common.flogger.FluentLogger; @@ -57,7 +59,6 @@ import google.registry.model.rde.RdeMode; 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,8 @@ import jakarta.inject.Singleton; import java.io.IOException; import java.io.Serializable; import java.time.Instant; +import java.util.List; +import java.util.NoSuchElementException; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.KvCoder; @@ -80,15 +83,19 @@ import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Distinct; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Filter; 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; @@ -96,6 +103,7 @@ import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.TypeDescriptor; +import org.hibernate.jpa.AvailableHints; /** * Definition of a Dataflow Flex template, which generates RDE/BRDA deposits. @@ -129,7 +137,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; *

{@link Domain}

* * After the most recent (live) domain resources are loaded from the corresponding history objects, - * we marshall them to deposit fragments and emit the (pending deposit: deposit fragment) pairs for + * we marshal them to deposit fragments and emit the (pending deposit: deposit fragment) pairs for * further processing. We also find all the hosts referenced by a given domain and emit pairs of * (host repo ID: pending deposit) for all RDE pending deposits for further processing. * @@ -143,11 +151,11 @@ import org.apache.beam.sdk.values.TypeDescriptor; * *

For subordinate hosts, we need to find the superordinate domain in order to properly handle * pending transfer in the deposit as well. So we first find the superordinate domain repo ID from - * the host and join the (superordinate domain repo ID: (subordinate host repo ID: (pending deposit: - * revision ID))) pair with the (domain repo ID: revision ID) pair obtained from the domain history - * query in order to map the host at watermark to the domain at watermark. We then proceed to create - * the (pending deposit: deposit fragment) pair for subordinate hosts using the added domain - * information. + * the host and join the (superordinate domain repo ID: (subordinate host repo ID: (pending + * deposits: revision IDs))) pair with the (domain repo ID: revision ID) pair obtained from the + * domain history query in order to map the host at watermark to the domain at watermark. We then + * proceed to create the (pending deposit: deposit fragment) pair for subordinate hosts using the + * added domain information. * *

Processing {@link DepositFragment}

* @@ -182,6 +190,25 @@ public class RdePipeline implements Serializable { private static final FluentLogger logger = FluentLogger.forEnclosingClass(); + // Metrics counters for tracking pipeline execution. + private static final Counter INCLUDED_REGISTRAR_COUNTER = + Metrics.counter("RDE", "IncludedRegistrar"); + private static final Counter REGISTRAR_FRAGMENT_COUNTER = + Metrics.counter("RDE", "RegistrarFragment"); + private static final Counter REFERENCED_HOST_COGBK_COUNTER = + Metrics.counter("RDE", "ReferencedHost"); + private static final Counter SUBORDINATE_HOST_COUNTER = Metrics.counter("RDE", "SubordinateHost"); + private static final Counter EXTERNAL_HOST_COUNTER = Metrics.counter("RDE", "ExternalHost"); + private static final Counter EXTERNAL_HOST_FRAGMENT_COUNTER = + Metrics.counter("RDE", "ExternalHostFragment"); + private static final Counter SUBORDINATE_HOST_FRAGMENT_COUNTER = + Metrics.counter("RDE", "SubordinateHostFragment"); + private static final Counter REFERENCED_SUBORDINATE_HOST_COUNTER = + Metrics.counter("RDE", "ReferencedSubordinateHost"); + private static final Counter ACTIVE_DOMAIN_COUNTER = Metrics.counter("RDE", "ActiveDomain"); + private static final Counter DOMAIN_FRAGMENT_COUNTER = Metrics.counter("RDE", "DomainFragment"); + private static final Counter REFERENCED_HOST_COUNTER = Metrics.counter("RDE", "ReferencedHost"); + @Inject RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils, CloudTasksUtils cloudTasksUtils) { this.options = options; @@ -262,12 +289,6 @@ public class RdePipeline implements Serializable { } private PCollection> processRegistrars(Pipeline pipeline) { - // Note that the namespace in the metric is not being used by Stackdriver, it just has to be - // non-empty. - // See: - // https://stackoverflow.com/questions/48530496/google-dataflow-custom-metrics-not-showing-on-stackdriver - Counter includedRegistrarCounter = Metrics.counter("RDE", "IncludedRegistrar"); - Counter registrarFragmentCounter = Metrics.counter("RDE", "RegistrarFragment"); return pipeline .apply( "Read all production Registrars", @@ -286,14 +307,14 @@ public class RdePipeline implements Serializable { .via( (String registrarRepoId) -> { VKey key = VKey.create(Registrar.class, registrarRepoId); - includedRegistrarCounter.inc(); + INCLUDED_REGISTRAR_COUNTER.inc(); Registrar registrar = tm().transact(() -> tm().loadByKey(key)); DepositFragment fragment = marshaller.marshalRegistrar(registrar); ImmutableSet> fragments = pendingDeposits.stream() .map(pending -> KV.of(pending, fragment)) .collect(toImmutableSet()); - registrarFragmentCounter.inc(fragments.size()); + REGISTRAR_FRAGMENT_COUNTER.inc(fragments.size()); return fragments; })); } @@ -303,42 +324,55 @@ public class RdePipeline implements Serializable { * *

Note that deleted and non-production resources are not included. * - * @return A KV pair of (repoId, revisionId), used to reconstruct the composite key for the + * @return A collection of (repoId -> revisionId) used to reconstruct the composite key for the * history entry. */ private PCollection> getMostRecentHistoryEntries( Pipeline pipeline, Class historyClass) { + String tldFilter = + historyClass == DomainHistory.class + ? " AND sub.resource.tld IN (SELECT id FROM Tld WHERE tldType = 'REAL')" + : ""; + String jpql = + String.format( + """ + SELECT repoId, revisionId FROM %1$s WHERE (repoId, modificationTime) IN ( + SELECT sub.repoId, MAX(sub.modificationTime) FROM %1$s sub + WHERE sub.modificationTime <= :watermark%2$s + GROUP BY sub.repoId + ) + AND resource.deletionTime > :watermark + AND COALESCE(resource.creationRegistrarId, '') NOT LIKE 'prober-%%' + AND COALESCE(resource.currentSponsorRegistrarId, '') NOT LIKE 'prober-%%' + AND COALESCE(resource.lastEppUpdateRegistrarId, '') NOT LIKE 'prober-%%' + """, + historyClass.getSimpleName(), tldFilter); return pipeline.apply( String.format("Load most recent %s", historyClass.getSimpleName()), RegistryJpaIO.read( - ("SELECT repoId, revisionId FROM %entity% WHERE (repoId, modificationTime) IN" - + " (SELECT repoId, MAX(modificationTime) FROM %entity% WHERE modificationTime" - + " <= :watermark GROUP BY repoId) AND resource.deletionTime > :watermark AND" - + " COALESCE(resource.creationRegistrarId, '') NOT LIKE 'prober-%' AND" - + " COALESCE(resource.currentSponsorRegistrarId, '') NOT LIKE 'prober-%' AND" - + " COALESCE(resource.lastEppUpdateRegistrarId, '') NOT LIKE 'prober-%' " - + (historyClass == DomainHistory.class - ? "AND resource.tld IN " + "(SELECT id FROM Tld WHERE tldType = 'REAL')" - : "")) - .replace("%entity%", historyClass.getSimpleName()), + jpql, ImmutableMap.of("watermark", watermark), Object[].class, row -> KV.of((String) row[0], (long) row[1])) .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 - // after CoGroupByKey (joining the revision IDs and the pending deposits on repo IDs), in - // #removedUnreferencedResources, duplicate revision IDs are sometimes introduced. Here we - // attempt to deduplicate the iterable. If it contains multiple revision IDs that are NOT the - // same, we have a more serious problem as we cannot be sure which one to use. We should use the - // highest revision ID, but we don't even know where it comes from, as the query should - // definitively only give us one revision ID per repo ID. In this case we have to abort and - // require manual intervention. + // The SQL query in getMostRecentHistoryEntries guarantees exactly one (repoId, revisionId) pair + // per entity. However, after multi-way joins via CoGroupByKey (e.g. when joining pending + // deposits or subordinate hosts on repoId), duplicate identical revision IDs can appear in + // the resulting Iterable. + // + // We deduplicate the iterable here. If it contains multiple revision IDs that are NOT + // identical, we have an illegal state because we cannot determine which historical revision is + // authoritative at the watermark. In that case, we abort and require manual intervention. + checkArgument( + !ids.isEmpty(), + "No revision IDs found for %s repo ID %s", + historyEntryClazz.getSimpleName(), + repoId); if (ids.size() != 1) { ImmutableSet dedupedIds = ImmutableSet.copyOf(ids); checkState( @@ -347,169 +381,246 @@ public class RdePipeline implements Serializable { historyEntryClazz.getSimpleName(), repoId, ids); - logger.atSevere().log( + logger.atInfo().log( "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(); + static + ImmutableMap loadResourcesByHistoryEntryIds( + Iterable> repoAndRevisionIds, + Class resourceClass, + Class historyEntryClass, + Instant watermark) { + ImmutableList> ids = ImmutableList.copyOf(repoAndRevisionIds); + if (ids.isEmpty()) { + return ImmutableMap.of(); + } + String[] repoIdArray = ids.stream().map(KV::getKey).toArray(String[]::new); + Long[] revisionIdArray = ids.stream().map(KV::getValue).toArray(Long[]::new); + String repoIdColumnName = + historyEntryClass.equals(DomainHistory.class) ? "domain_repo_id" : "host_repo_id"; + // Unfortunately Hibernate doesn't play nice with selecting by composite primary keys. We cannot + // directly say "WHERE (repoId, revisionId) IN (repoIdAndRevisionIdPairs)" in any way in HQL. + // As a result, we must use the native query format to quickly select against the (repoId, + // revisionId) primary key index. Just make sure not to use batch sizes in the tens of thousands + // (default is 500), otherwise the query could get too long. + String nativeQuerySql = + String.format( + """ + SELECT * FROM "%s" WHERE (%s, history_revision_id) IN ( + SELECT * FROM UNNEST(:repoIds\\:\\:text[], :revisionIds\\:\\:bigint[])) + """, + historyEntryClass.getSimpleName(), repoIdColumnName); + ImmutableMap result = + tm().transact( + () -> { + @SuppressWarnings("unchecked") + List queryResult = + tm().getEntityManager() + .createNativeQuery(nativeQuerySql, historyEntryClass) + .setParameter("repoIds", repoIdArray) + .setParameter("revisionIds", revisionIdArray) + .setHint(AvailableHints.HINT_READ_ONLY, true) + .getResultList(); + // Flush the context so we can GC aggressively + tm().getEntityManager().clear(); + return queryResult.stream() + .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()) { + ImmutableSet expectedRepoIds = ids.stream().map(KV::getKey).collect(toImmutableSet()); + throw new NoSuchElementException( + String.format( + "Expected to find the following %s history entries but they were missing: %s", + historyEntryClass.getSimpleName(), + Sets.difference(expectedRepoIds, 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 - * to be included in the corresponding pending deposit. + *

The (repoId, pendingDeposit) pairs denote hosts that are referenced from a domain (built up + * when processing domains earlier). We essentially want to filter out the hostHistories to only + * contain these hosts. * - *

The (repoId, revisionId) pairs come from the most recent history entry query, which can be - * used to load the embedded resources themselves. - * - * @return a pair of (repoId, ([pendingDeposit], [revisionId])) where neither the pendingDeposit - * nor the revisionId list is empty. + * @return a collection of (repoId -> (pending deposits, revisionId)) where neither the + * pendingDeposit nor the revisionId list is empty. */ - 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); - return KeyedPCollectionTuple.of(PENDING_DEPOSIT, referencedResources) - .and(REVISION_ID, historyEntries) + private static PCollection> removeUnreferencedHosts( + PCollection> referencedHosts, + PCollection> hostHistories) { + PCollection> uniqueHosts = + referencedHosts + .setCoder(KvCoder.of(StringUtf8Coder.of(), PendingDepositCoder.of())) + .apply("Deduplicate hosts for grouping", Distinct.create()); + return KeyedPCollectionTuple.of(PENDING_DEPOSIT, uniqueHosts) + .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 Hosts", Filter.by( (KV kv) -> { boolean toInclude = - // If a resource does not have corresponding pending deposit, it is not - // referenced and should not be included. - kv.getValue().getAll(PENDING_DEPOSIT).iterator().hasNext() - // If a resource does not have revision id (this should not happen, as - // every referenced resource must be valid at watermark time, therefore + // If a host does not have corresponding pending deposit, it is not referenced + // and should not be included. + !Iterables.isEmpty(kv.getValue().getAll(PENDING_DEPOSIT)) + // If a host does not have revision id (this should not happen, as + // every referenced host must be valid at watermark time, therefore // be embedded in a history entry valid at watermark time, otherwise // the domain cannot reference it), there is no way for us to find the - // history entry and load the embedded resource. So we ignore the resource + // history entry and load the embedded host. So we ignore the host // to keep the downstream process simple. - && kv.getValue().getAll(REVISION_ID).iterator().hasNext(); + && !Iterables.isEmpty(kv.getValue().getAll(REVISION_ID)); if (toInclude) { - referencedResourceCounter.inc(); + REFERENCED_HOST_COGBK_COUNTER.inc(); } return toInclude; })); } private PCollectionTuple processDomainHistories(PCollection> domainHistories) { - 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.getHistoryEntryLoadBatchSize(); + int numShards = options.getNumHistoryEntryShards(); + return domainHistories + .apply( + // Batching only combines elements with the same key, so we need to shard + "Split domain histories across shards for batched retrieval", + WithKeys.>of( + kv -> Math.floorMod(kv.getKey().hashCode(), numShards)) + .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) { + loadResourcesByHistoryEntryIds( + element.getValue(), Domain.class, DomainHistory.class, watermark) + .values() + .forEach(d -> processSingleDomain(d, receiver)); + } + }) + .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); + } + + private void processSingleDomain(Domain domain, DoFn.MultiOutputReceiver receiver) { + ACTIVE_DOMAIN_COUNTER.inc(); + pendingDeposits.stream() + .filter(pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) + .forEach( + pendingDeposit -> { + DOMAIN_FRAGMENT_COUNTER.inc(); + receiver + .get(DOMAIN_FRAGMENTS) + .output( + KV.of( + pendingDeposit, marshaller.marshalDomain(domain, pendingDeposit.mode()))); + + if (pendingDeposit.mode() == RdeMode.FULL && domain.getNsHosts() != null) { + REFERENCED_HOST_COUNTER.inc(domain.getNsHosts().size()); + domain + .getNsHosts() + .forEach( + hostKey -> + receiver + .get(REFERENCED_HOSTS) + .output(KV.of((String) hostKey.getKey(), pendingDeposit))); + } + }); } private PCollectionTuple processHostHistories( PCollection> referencedHosts, PCollection> hostHistories) { - 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.getHistoryEntryLoadBatchSize(); + int numShards = options.getNumHistoryEntryShards(); + return removeUnreferencedHosts(referencedHosts, hostHistories) .apply( - "Map external DomainResource to DepositFragment and process subordinate domains", + // Batching only combines elements with the same key, so we need to shard + "Split host histories across shards for batched retrieval", + WithKeys.>of( + kv -> Math.floorMod(kv.getKey().hashCode(), numShards)) + .withKeyType(integers())) + .apply( + "Group referenced hosts into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) + .apply( + "Map external Host to DepositFragment and route subordinate hosts", 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) { + ImmutableList> batchElements = + ImmutableList.copyOf(element.getValue()); + ImmutableSet> hostKeys = + batchElements.stream() + .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, watermark); + for (KV kv : batchElements) { + 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()) { + SUBORDINATE_HOST_COUNTER.inc(); + receiver + .get(SUPERORDINATE_DOMAINS) + .output( + // The output are pairs of (superordinateDomainRepoId, + // (subordinateHostRepoId, (pendingDeposits, revisionIds))). + KV.of((String) host.getSuperordinateDomain().getKey(), kv)); + } else { + // We can just directly marshal and send out external hosts + EXTERNAL_HOST_COUNTER.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 -> { + EXTERNAL_HOST_FRAGMENT_COUNTER.inc(); + receiver + .get(EXTERNAL_HOST_FRAGMENTS) + .output(KV.of(pendingDeposit, fragment)); + }); + } } } }) @@ -521,7 +632,7 @@ public class RdePipeline implements Serializable { * obtained from its superordinate domain. * * @param superordinateDomains Pairs of (superordinateDomainRepoId, (subordinateHostRepoId, - * (pendingDeposit, revisionId))). This collection maps the subordinate host and the pending + * (pendingDeposits, revisionIds))). This collection maps the subordinate host and the pending * deposit to include it to its superordinate domain. * @param domainHistories Pairs of (domainRepoId, revisionId). This collection helps us find the * historical superordinate domain from its history entry and is obtained from calling {@link @@ -530,59 +641,98 @@ public class RdePipeline implements Serializable { private PCollection> processSubordinateHosts( PCollection>> superordinateDomains, PCollection> domainHistories) { - Counter subordinateHostFragmentCounter = Metrics.counter("RDE", "SubordinateHostFragment"); - Counter referencedSubordinateHostCounter = Metrics.counter("RDE", "ReferencedSubordinateHost"); - return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT, superordinateDomains) + int batchSize = options.getHistoryEntryLoadBatchSize(); + int numShards = options.getNumHistoryEntryShards(); + return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID, superordinateDomains) .and(REVISION_ID, domainHistories) .apply("Join Host:PendingDeposits with DomainHistory on Domain", CoGroupByKey.create()) .apply( - " Remove unreferenced Domain", + "Remove Domains without subordinate hosts", Filter.by( kv -> { boolean toInclude = - kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT).iterator().hasNext() - && kv.getValue().getAll(REVISION_ID).iterator().hasNext(); + !Iterables.isEmpty( + kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID)) + && !Iterables.isEmpty(kv.getValue().getAll(REVISION_ID)); if (toInclude) { - referencedSubordinateHostCounter.inc(); + REFERENCED_SUBORDINATE_HOST_COUNTER.inc(); } return toInclude; })) + .apply( + // Batching only combines elements with the same key, so we need to shard + "Split superordinate domains across shards for batched retrieval", + WithKeys.>of( + kv -> Math.floorMod(kv.getKey().hashCode(), numShards)) + .withKeyType(integers())) + .apply( + "Group superordinate domains 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) { + ImmutableList> batchElements = + ImmutableList.copyOf(element.getValue()); + ImmutableSet> domainKeys = + batchElements.stream() + .map( + kv -> + KV.of( + kv.getKey(), + getSingleRevisionId( + DomainHistory.class, + kv.getKey(), + kv.getValue().getAll(REVISION_ID)))) + .collect(toImmutableSet()); + ImmutableSet> hostKeys = + batchElements.stream() + .flatMap( + kv -> + Streams.stream( + kv.getValue() + .getAll(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID)) + .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, watermark); + ImmutableMap loadedHosts = + loadResourcesByHistoryEntryIds( + hostKeys, Host.class, HostHistory.class, watermark); + for (KV kv : batchElements) { + 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)); + kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID)) { + Host host = loadedHosts.get(hostToPendingDeposits.getKey()); DepositFragment fragment = marshaller.marshalSubordinateHost(host, superordinateDomain); Streams.stream(hostToPendingDeposits.getValue().getAll(PENDING_DEPOSIT)) .distinct() .forEach( pendingDeposit -> { - subordinateHostFragmentCounter.inc(); - results.add(KV.of(pendingDeposit, fragment)); + SUBORDINATE_HOST_FRAGMENT_COUNTER.inc(); + receiver.output(KV.of(pendingDeposit, fragment)); }); } - return results.build(); - })); + } + } + })); } /** @@ -634,8 +784,8 @@ public class RdePipeline implements Serializable { protected static final TupleTag PENDING_DEPOSIT = new TupleTag<>() {}; - protected static final TupleTag> HOST_TO_PENDING_DEPOSIT = - new TupleTag<>() {}; + protected static final TupleTag> + HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID = new TupleTag<>() {}; protected static final TupleTag REVISION_ID = new TupleTag<>() {}; } 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 c365494b6..4627c8de7 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,17 @@ 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 getHistoryEntryLoadBatchSize(); + + void setHistoryEntryLoadBatchSize(int value); + + @Description("The number of shards to use when splitting items into batches") + @Default.Integer(1000) + int getNumHistoryEntryShards(); + + void setNumHistoryEntryShards(int value); } diff --git a/core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json b/core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json index 3b39c5725..f2db912bf 100644 --- a/core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json +++ b/core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json @@ -43,6 +43,24 @@ "regexes": [ "[A-Za-z0-9\\-_]+" ] + }, + { + "name": "historyEntryLoadBatchSize", + "label": "History entry load batch size.", + "helpText": "The number of history entries to load from the database in one operation.", + "is_optional": true, + "regexes": [ + "^[1-9][0-9]*$" + ] + }, + { + "name": "numHistoryEntryShards", + "label": "Number of history entry shards.", + "helpText": "The number of shards to split across when batching history entries.", + "is_optional": true, + "regexes": [ + "^[1-9][0-9]*$" + ] } ] } 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 bd8fa1314..2c6d176e6 100644 --- a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java +++ b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java @@ -29,6 +29,7 @@ import static google.registry.rde.RdeResourceType.DOMAIN; import static google.registry.rde.RdeResourceType.HOST; import static google.registry.rde.RdeResourceType.REGISTRAR; import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.loadByEntity; import static google.registry.testing.DatabaseHelper.newDomain; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistActiveHost; @@ -85,6 +86,7 @@ import google.registry.testing.FakeKeyringModule; 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 +168,42 @@ public class RdePipelineTest { .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 @@ -252,12 +260,12 @@ public class RdePipelineTest { .build()); persistDomainHistory(kittyDomain); // Should not appear because the TLD is not included in a pending deposit. - persistDomainHistory(persistEppResource(newDomain("lol.cat"))); + persistDomainHistory(persistActiveDomain("lol.cat")); // To be deleted. 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 +433,81 @@ public class RdePipelineTest { pipeline.run().waitUntilFinish(); } + @Test + void testSuccess_createFragments_smallBatchSize() { + options.setHistoryEntryLoadBatchSize(1); + testSuccess_createFragments(); + } + + @Test + void testSuccess_createFragments_multiBatch() { + options.setHistoryEntryLoadBatchSize(2); + testSuccess_createFragments(); + } + + @Test + void testFailure_missingHistoryEntry() { + NoSuchElementException thrown = + assertThrows( + NoSuchElementException.class, + () -> + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of(KV.of("nonexistent", 12345L)), + Domain.class, + DomainHistory.class, + clock.now())); + assertThat(thrown).hasMessageThat().contains("nonexistent"); + } + + @Test + void testSuccess_loadResourcesByHistoryEntryIds_multipleRevisions() { + Domain domain = loadByEntity(persistActiveDomain("multirev.soy")); + DomainHistory history1 = persistDomainHistory(domain); + clock.advanceOneMilli(); + Domain updatedDomain = + domain.asBuilder().setPersistedCurrentSponsorRegistrarId("NewRegistrar").build(); + DomainHistory history2 = persistDomainHistory(updatedDomain); + + // Verify loading specific revision 1 returns history1 entity + ImmutableMap loaded1 = + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of(KV.of(domain.getRepoId(), history1.getRevisionId())), + Domain.class, + DomainHistory.class, + now); + assertThat(loaded1.get(domain.getRepoId()).getCurrentSponsorRegistrarId()) + .isEqualTo("TheRegistrar"); + + // Verify loading specific revision 2 returns history2 entity + ImmutableMap loaded2 = + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of(KV.of(domain.getRepoId(), history2.getRevisionId())), + Domain.class, + DomainHistory.class, + now); + assertThat(loaded2.get(domain.getRepoId()).getCurrentSponsorRegistrarId()) + .isEqualTo("NewRegistrar"); + } + + @Test + void testSuccess_loadResourcesByHistoryEntryIds_batchMultipleEntities() { + Domain domain1 = persistActiveDomain("batch1.soy"); + DomainHistory history1 = persistDomainHistory(domain1); + Domain domain2 = persistActiveDomain("batch2.soy"); + DomainHistory history2 = persistDomainHistory(domain2); + + ImmutableMap loaded = + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of( + KV.of(domain1.getRepoId(), history1.getRevisionId()), + KV.of(domain2.getRepoId(), history2.getRevisionId())), + Domain.class, + DomainHistory.class, + now); + + assertThat(loaded.keySet()).containsExactly(domain1.getRepoId(), domain2.getRepoId()); + } + // The GCS folder listing can be a bit flaky, so retry if necessary @RetryingTest(4) void testSuccess_persistData() throws Exception {