mirrors/nomulus
1
0
Fork 0
mirror of https://github.com/google/nomulus synced 2026-09-29 11:15:45 +00:00
Code Issues Releases Wiki Activity

Batch database queries in RdePipeline (#3190)

Browse Source 
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<String, Long> 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
This commit is contained in:
gbrodman
2026-08-15 01:20:27 +00:00
committed by GitHub
parent 1423474fb1
commit a8f01d590d
4 changed files with 497 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/main/java/google/registry/beam/rde/RdePipeline.java
+352 -202
View File
@@ -16,25 +16,27 @@ package google.registry.beam.rde;
import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState; 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 com.google.common.collect.ImmutableSet.toImmutableSet;
import static google.registry.beam.rde.RdePipeline.TupleTags.DOMAIN_FRAGMENTS; 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.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.PENDING_DEPOSIT;
import static google.registry.beam.rde.RdePipeline.TupleTags.REFERENCED_HOSTS; 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.REVISION_ID;
import static google.registry.beam.rde.RdePipeline.TupleTags.SUPERORDINATE_DOMAINS; 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.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.util.SafeSerializationUtils.safeDeserializeCollection; import static google.registry.util.SafeSerializationUtils.safeDeserializeCollection;
import static google.registry.util.SafeSerializationUtils.serializeCollection; import static google.registry.util.SafeSerializationUtils.serializeCollection;
import static google.registry.util.SerializeUtils.decodeBase64; import static google.registry.util.SerializeUtils.decodeBase64;
import static google.registry.util.SerializeUtils.encodeBase64; 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 static org.apache.beam.sdk.values.TypeDescriptors.kvs;
import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets; import com.google.common.collect.Sets;
import com.google.common.collect.Streams; import com.google.common.collect.Streams;
import com.google.common.flogger.FluentLogger; 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;
import google.registry.model.registrar.Registrar.Type; import google.registry.model.registrar.Registrar.Type;
import google.registry.model.reporting.HistoryEntry; import google.registry.model.reporting.HistoryEntry;
import google.registry.model.reporting.HistoryEntry.HistoryEntryId;
import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
import google.registry.persistence.VKey; import google.registry.persistence.VKey;
import google.registry.rde.DepositFragment; import google.registry.rde.DepositFragment;
@@ -71,6 +72,8 @@ import jakarta.inject.Singleton;
import java.io.IOException; import java.io.IOException;
import java.io.Serializable; import java.io.Serializable;
import java.time.Instant; import java.time.Instant;
import java.util.List;
import java.util.NoSuchElementException;
import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.KvCoder; 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.Counter;
import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.PipelineOptionsFactory; 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.DoFn;
import org.apache.beam.sdk.transforms.Filter; import org.apache.beam.sdk.transforms.Filter;
import org.apache.beam.sdk.transforms.FlatMapElements; import org.apache.beam.sdk.transforms.FlatMapElements;
import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.GroupByKey; 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.ParDo;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.join.CoGbkResult; import org.apache.beam.sdk.transforms.join.CoGbkResult;
import org.apache.beam.sdk.transforms.join.CoGroupByKey; import org.apache.beam.sdk.transforms.join.CoGroupByKey;
import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple; 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.KV;
import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList; 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.TupleTag;
import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptor;
import org.hibernate.jpa.AvailableHints;
/** /**
* Definition of a Dataflow Flex template, which generates RDE/BRDA deposits. * Definition of a Dataflow Flex template, which generates RDE/BRDA deposits.
@@ -129,7 +137,7 @@ import org.apache.beam.sdk.values.TypeDescriptor;
* <h3>{@link Domain}</h3> * <h3>{@link Domain}</h3>
* *
* After the most recent (live) domain resources are loaded from the corresponding history objects, * 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 * 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. * (host repo ID: pending deposit) for all RDE pending deposits for further processing.
* *
@@ -143,11 +151,11 @@ import org.apache.beam.sdk.values.TypeDescriptor;
* *
* <p>For subordinate hosts, we need to find the superordinate domain in order to properly handle * <p>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 * 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: * the host and join the (superordinate domain repo ID: (subordinate host repo ID: (pending
* revision ID))) pair with the (domain repo ID: revision ID) pair obtained from the domain history * deposits: revision IDs))) pair with the (domain repo ID: revision ID) pair obtained from the
* query in order to map the host at watermark to the domain at watermark. We then proceed to create * domain history query in order to map the host at watermark to the domain at watermark. We then
* the (pending deposit: deposit fragment) pair for subordinate hosts using the added domain * proceed to create the (pending deposit: deposit fragment) pair for subordinate hosts using the
* information. * added domain information.
* *
* <h2>Processing {@link DepositFragment}</h2> * <h2>Processing {@link DepositFragment}</h2>
* *
@@ -182,6 +190,25 @@ public class RdePipeline implements Serializable {
private static final FluentLogger logger = FluentLogger.forEnclosingClass(); 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 @Inject
RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils, CloudTasksUtils cloudTasksUtils) { RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils, CloudTasksUtils cloudTasksUtils) {
this.options = options; this.options = options;
@@ -262,12 +289,6 @@ public class RdePipeline implements Serializable {
} }
private PCollection<KV<PendingDeposit, DepositFragment>> processRegistrars(Pipeline pipeline) { private PCollection<KV<PendingDeposit, DepositFragment>> 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 return pipeline
.apply( .apply(
"Read all production Registrars", "Read all production Registrars",
@@ -286,14 +307,14 @@ public class RdePipeline implements Serializable {
.via( .via(
(String registrarRepoId) -> { (String registrarRepoId) -> {
VKey<Registrar> key = VKey.create(Registrar.class, registrarRepoId); VKey<Registrar> key = VKey.create(Registrar.class, registrarRepoId);
includedRegistrarCounter.inc(); INCLUDED_REGISTRAR_COUNTER.inc();
Registrar registrar = tm().transact(() -> tm().loadByKey(key)); Registrar registrar = tm().transact(() -> tm().loadByKey(key));
DepositFragment fragment = marshaller.marshalRegistrar(registrar); DepositFragment fragment = marshaller.marshalRegistrar(registrar);
ImmutableSet<KV<PendingDeposit, DepositFragment>> fragments = ImmutableSet<KV<PendingDeposit, DepositFragment>> fragments =
pendingDeposits.stream() pendingDeposits.stream()
.map(pending -> KV.of(pending, fragment)) .map(pending -> KV.of(pending, fragment))
.collect(toImmutableSet()); .collect(toImmutableSet());
registrarFragmentCounter.inc(fragments.size()); REGISTRAR_FRAGMENT_COUNTER.inc(fragments.size());
return fragments; return fragments;
})); }));
} }
@@ -303,42 +324,55 @@ public class RdePipeline implements Serializable {
* *
* <p>Note that deleted and non-production resources are not included. * <p>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. * history entry.
*/ */
private <T extends HistoryEntry> PCollection<KV<String, Long>> getMostRecentHistoryEntries( private <T extends HistoryEntry> PCollection<KV<String, Long>> getMostRecentHistoryEntries(
Pipeline pipeline, Class<T> historyClass) { Pipeline pipeline, Class<T> 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( return pipeline.apply(
String.format("Load most recent %s", historyClass.getSimpleName()), String.format("Load most recent %s", historyClass.getSimpleName()),
RegistryJpaIO.read( RegistryJpaIO.read(
("SELECT repoId, revisionId FROM %entity% WHERE (repoId, modificationTime) IN" jpql,
+ " (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()),
ImmutableMap.of("watermark", watermark), ImmutableMap.of("watermark", watermark),
Object[].class, Object[].class,
row -> KV.of((String) row[0], (long) row[1])) row -> KV.of((String) row[0], (long) row[1]))
.withCoder(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))); .withCoder(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of())));
} }
private <T extends HistoryEntry> EppResource loadResourceByHistoryEntryId( private static long getSingleRevisionId(
Class<T> historyEntryClazz, String repoId, Iterable<Long> revisionIds) { Class<? extends HistoryEntry> historyEntryClazz, String repoId, Iterable<Long> revisionIds) {
ImmutableList<Long> ids = ImmutableList.copyOf(revisionIds); ImmutableList<Long> ids = ImmutableList.copyOf(revisionIds);
// The size should always be 1 because we are only getting one repo ID -> revision ID pair per // The SQL query in getMostRecentHistoryEntries guarantees exactly one (repoId, revisionId) pair
// repo ID from the source transform (the JPA query in the method above). But for some reason // per entity. However, after multi-way joins via CoGroupByKey (e.g. when joining pending
// after CoGroupByKey (joining the revision IDs and the pending deposits on repo IDs), in // deposits or subordinate hosts on repoId), duplicate identical revision IDs can appear in
// #removedUnreferencedResources, duplicate revision IDs are sometimes introduced. Here we // the resulting Iterable<Long>.
// 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 // We deduplicate the iterable here. If it contains multiple revision IDs that are NOT
// highest revision ID, but we don't even know where it comes from, as the query should // identical, we have an illegal state because we cannot determine which historical revision is
// definitively only give us one revision ID per repo ID. In this case we have to abort and // authoritative at the watermark. In that case, we abort and require manual intervention.
// require manual intervention. checkArgument(
!ids.isEmpty(),
"No revision IDs found for %s repo ID %s",
historyEntryClazz.getSimpleName(),
repoId);
if (ids.size() != 1) { if (ids.size() != 1) {
ImmutableSet<Long> dedupedIds = ImmutableSet.copyOf(ids); ImmutableSet<Long> dedupedIds = ImmutableSet.copyOf(ids);
checkState( checkState(
@@ -347,169 +381,246 @@ public class RdePipeline implements Serializable {
historyEntryClazz.getSimpleName(), historyEntryClazz.getSimpleName(),
repoId, repoId,
ids); ids);
logger.atSevere().log( logger.atInfo().log(
"Duplicate revision IDs detected for %s repo ID %s: %s", "Duplicate revision IDs detected for %s repo ID %s: %s",
historyEntryClazz.getSimpleName(), repoId, ids); historyEntryClazz.getSimpleName(), repoId, ids);
} }
return loadResourceByHistoryEntryId(historyEntryClazz, repoId, ids.get(0)); return ids.getFirst();
} }
private <T extends HistoryEntry> EppResource loadResourceByHistoryEntryId( static <E extends EppResource, H extends HistoryEntry>
Class<T> historyEntryClazz, String repoId, long revisionId) { ImmutableMap<String, E> loadResourcesByHistoryEntryIds(
return tm().transact( Iterable<KV<String, Long>> repoAndRevisionIds,
() -> Class<E> resourceClass,
tm().loadByKey( Class<H> historyEntryClass,
VKey.create(historyEntryClazz, new HistoryEntryId(repoId, revisionId)))) Instant watermark) {
.getResourceAtPointInTime() ImmutableList<KV<String, Long>> ids = ImmutableList.copyOf(repoAndRevisionIds);
.map(resource -> resource.cloneProjectedAtTime(watermark)) if (ids.isEmpty()) {
.get(); 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<String, E> result =
tm().transact(
() -> {
@SuppressWarnings("unchecked")
List<H> 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<String> 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. * revisionId) on the repoId.
* *
* <p>The (repoId, pendingDeposit) pairs denote hosts that are referenced from a domain, that are * <p>The (repoId, pendingDeposit) pairs denote hosts that are referenced from a domain (built up
* to be included in the corresponding pending deposit. * when processing domains earlier). We essentially want to filter out the hostHistories to only
* contain these hosts.
* *
* <p>The (repoId, revisionId) pairs come from the most recent history entry query, which can be * @return a collection of (repoId -> (pending deposits, revisionId)) where neither the
* used to load the embedded resources themselves. * pendingDeposit nor the revisionId list is empty.
*
* @return a pair of (repoId, ([pendingDeposit], [revisionId])) where neither the pendingDeposit
* nor the revisionId list is empty.
*/ */
private static PCollection<KV<String, CoGbkResult>> removeUnreferencedResource( private static PCollection<KV<String, CoGbkResult>> removeUnreferencedHosts(
PCollection<KV<String, PendingDeposit>> referencedResources, PCollection<KV<String, PendingDeposit>> referencedHosts,
PCollection<KV<String, Long>> historyEntries, PCollection<KV<String, Long>> hostHistories) {
Class<? extends EppResource> resourceClazz) { PCollection<KV<String, PendingDeposit>> uniqueHosts =
String resourceName = resourceClazz.getSimpleName(); referencedHosts
Class<? extends HistoryEntry> historyEntryClazz = .setCoder(KvCoder.of(StringUtf8Coder.of(), PendingDepositCoder.of()))
RESOURCE_TYPES_TO_HISTORY_TYPES.get(resourceClazz); .apply("Deduplicate hosts for grouping", Distinct.create());
String historyEntryName = historyEntryClazz.getSimpleName(); return KeyedPCollectionTuple.of(PENDING_DEPOSIT, uniqueHosts)
Counter referencedResourceCounter = Metrics.counter("RDE", "Referenced" + resourceName); .and(REVISION_ID, hostHistories)
return KeyedPCollectionTuple.of(PENDING_DEPOSIT, referencedResources) .apply("Join PendingDeposit with HostHistory revision ID on Host", CoGroupByKey.create())
.and(REVISION_ID, historyEntries)
.apply( .apply(
String.format( "Remove unreferenced Hosts",
"Join PendingDeposit with %s revision ID on %s", historyEntryName, resourceName),
CoGroupByKey.create())
.apply(
String.format("Remove unreferenced %s", resourceName),
Filter.by( Filter.by(
(KV<String, CoGbkResult> kv) -> { (KV<String, CoGbkResult> kv) -> {
boolean toInclude = boolean toInclude =
// If a resource does not have corresponding pending deposit, it is not // If a host does not have corresponding pending deposit, it is not referenced
// referenced and should not be included. // and should not be included.
kv.getValue().getAll(PENDING_DEPOSIT).iterator().hasNext() !Iterables.isEmpty(kv.getValue().getAll(PENDING_DEPOSIT))
// If a resource does not have revision id (this should not happen, as // If a host does not have revision id (this should not happen, as
// every referenced resource must be valid at watermark time, therefore // every referenced host must be valid at watermark time, therefore
// be embedded in a history entry valid at watermark time, otherwise // 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 // 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. // to keep the downstream process simple.
&& kv.getValue().getAll(REVISION_ID).iterator().hasNext(); && !Iterables.isEmpty(kv.getValue().getAll(REVISION_ID));
if (toInclude) { if (toInclude) {
referencedResourceCounter.inc(); REFERENCED_HOST_COGBK_COUNTER.inc();
} }
return toInclude; return toInclude;
})); }));
} }
private PCollectionTuple processDomainHistories(PCollection<KV<String, Long>> domainHistories) { private PCollectionTuple processDomainHistories(PCollection<KV<String, Long>> domainHistories) {
Counter activeDomainCounter = Metrics.counter("RDE", "ActiveDomainBase"); int batchSize = options.getHistoryEntryLoadBatchSize();
Counter domainFragmentCounter = Metrics.counter("RDE", "DomainFragment"); int numShards = options.getNumHistoryEntryShards();
Counter referencedHostCounter = Metrics.counter("RDE", "ReferencedHost"); return domainHistories
return domainHistories.apply( .apply(
"Map DomainHistory to DepositFragment and emit referenced Host", // Batching only combines elements with the same key, so we need to shard
ParDo.of( "Split domain histories across shards for batched retrieval",
new DoFn<KV<String, Long>, KV<PendingDeposit, DepositFragment>>() { WithKeys.<Integer, KV<String, Long>>of(
@ProcessElement kv -> Math.floorMod(kv.getKey().hashCode(), numShards))
public void processElement( .withKeyType(integers()))
@Element KV<String, Long> kv, MultiOutputReceiver receiver) { .apply(
activeDomainCounter.inc(); "Group domain histories into batches",
Domain domain = GroupIntoBatches.<Integer, KV<String, Long>>ofSize(batchSize).withShardedKey())
(Domain) .apply(
loadResourceByHistoryEntryId( "Map DomainHistory to DepositFragment and emit referenced Host",
DomainHistory.class, kv.getKey(), kv.getValue()); ParDo.of(
pendingDeposits.stream() new DoFn<
.filter(pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) KV<ShardedKey<Integer>, Iterable<KV<String, Long>>>,
.forEach( KV<PendingDeposit, DepositFragment>>() {
pendingDeposit -> { @ProcessElement
// Domains are always deposited in both modes. public void processElement(
domainFragmentCounter.inc(); @Element KV<ShardedKey<Integer>, Iterable<KV<String, Long>>> element,
receiver MultiOutputReceiver receiver) {
.get(DOMAIN_FRAGMENTS) loadResourcesByHistoryEntryIds(
.output( element.getValue(), Domain.class, DomainHistory.class, watermark)
KV.of( .values()
pendingDeposit, .forEach(d -> processSingleDomain(d, receiver));
marshaller.marshalDomain(domain, pendingDeposit.mode()))); }
// Hosts are only deposited in RDE, not BRDA. })
if (pendingDeposit.mode() == RdeMode.FULL) { .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS)));
if (domain.getNsHosts() != null) { }
referencedHostCounter.inc(domain.getNsHosts().size());
domain private void processSingleDomain(Domain domain, DoFn.MultiOutputReceiver receiver) {
.getNsHosts() ACTIVE_DOMAIN_COUNTER.inc();
.forEach( pendingDeposits.stream()
hostKey -> .filter(pendingDeposit -> pendingDeposit.tld().equals(domain.getTld()))
receiver .forEach(
.get(REFERENCED_HOSTS) pendingDeposit -> {
.output( DOMAIN_FRAGMENT_COUNTER.inc();
KV.of( receiver
(String) hostKey.getKey(), .get(DOMAIN_FRAGMENTS)
pendingDeposit))); .output(
} KV.of(
} pendingDeposit, marshaller.marshalDomain(domain, pendingDeposit.mode())));
});
} if (pendingDeposit.mode() == RdeMode.FULL && domain.getNsHosts() != null) {
}) REFERENCED_HOST_COUNTER.inc(domain.getNsHosts().size());
.withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); domain
.getNsHosts()
.forEach(
hostKey ->
receiver
.get(REFERENCED_HOSTS)
.output(KV.of((String) hostKey.getKey(), pendingDeposit)));
}
});
} }
private PCollectionTuple processHostHistories( private PCollectionTuple processHostHistories(
PCollection<KV<String, PendingDeposit>> referencedHosts, PCollection<KV<String, PendingDeposit>> referencedHosts,
PCollection<KV<String, Long>> hostHistories) { PCollection<KV<String, Long>> hostHistories) {
Counter subordinateHostCounter = Metrics.counter("RDE", "SubordinateHost"); int batchSize = options.getHistoryEntryLoadBatchSize();
Counter externalHostCounter = Metrics.counter("RDE", "ExternalHost"); int numShards = options.getNumHistoryEntryShards();
Counter externalHostFragmentCounter = Metrics.counter("RDE", "ExternalHostFragment"); return removeUnreferencedHosts(referencedHosts, hostHistories)
return removeUnreferencedResource(referencedHosts, hostHistories, Host.class)
.apply( .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.<Integer, KV<String, CoGbkResult>>of(
kv -> Math.floorMod(kv.getKey().hashCode(), numShards))
.withKeyType(integers()))
.apply(
"Group referenced hosts into batches",
GroupIntoBatches.<Integer, KV<String, CoGbkResult>>ofSize(batchSize).withShardedKey())
.apply(
"Map external Host to DepositFragment and route subordinate hosts",
ParDo.of( ParDo.of(
new DoFn<KV<String, CoGbkResult>, KV<PendingDeposit, DepositFragment>>() { new DoFn<
KV<ShardedKey<Integer>, Iterable<KV<String, CoGbkResult>>>,
KV<PendingDeposit, DepositFragment>>() {
@ProcessElement @ProcessElement
public void processElement( public void processElement(
@Element KV<String, CoGbkResult> kv, MultiOutputReceiver receiver) { @Element
Host host = KV<ShardedKey<Integer>, Iterable<KV<String, CoGbkResult>>> element,
(Host) MultiOutputReceiver receiver) {
loadResourceByHistoryEntryId( ImmutableList<KV<String, CoGbkResult>> batchElements =
HostHistory.class, ImmutableList.copyOf(element.getValue());
kv.getKey(), ImmutableSet<KV<String, Long>> hostKeys =
kv.getValue().getAll(REVISION_ID)); batchElements.stream()
// When a host is subordinate, we need to find its superordinate domain and .map(
// include it in the deposit as well. kv ->
if (host.isSubordinate()) { KV.of(
subordinateHostCounter.inc(); kv.getKey(),
receiver getSingleRevisionId(
.get(SUPERORDINATE_DOMAINS) HostHistory.class,
.output( kv.getKey(),
// The output are pairs of kv.getValue().getAll(REVISION_ID))))
// (superordinateDomainRepoId, .collect(toImmutableSet());
// (subordinateHostRepoId, (pendingDeposit, revisionId))). ImmutableMap<String, Host> loadedHosts =
KV.of((String) host.getSuperordinateDomain().getKey(), kv)); loadResourcesByHistoryEntryIds(
} else { hostKeys, Host.class, HostHistory.class, watermark);
externalHostCounter.inc(); for (KV<String, CoGbkResult> kv : batchElements) {
DepositFragment fragment = marshaller.marshalExternalHost(host); Host host = loadedHosts.get(kv.getKey());
Streams.stream(kv.getValue().getAll(PENDING_DEPOSIT)) // When a host is subordinate, we need to find its superordinate domain
// The same host could be used by multiple domains, therefore // and include it in the deposit as well.
// matched to the same pending deposit multiple times. if (host.isSubordinate()) {
.distinct() SUBORDINATE_HOST_COUNTER.inc();
.forEach( receiver
pendingDeposit -> { .get(SUPERORDINATE_DOMAINS)
externalHostFragmentCounter.inc(); .output(
receiver // The output are pairs of (superordinateDomainRepoId,
.get(EXTERNAL_HOST_FRAGMENTS) // (subordinateHostRepoId, (pendingDeposits, revisionIds))).
.output(KV.of(pendingDeposit, fragment)); 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. * obtained from its superordinate domain.
* *
* @param superordinateDomains Pairs of (superordinateDomainRepoId, (subordinateHostRepoId, * @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. * deposit to include it to its superordinate domain.
* @param domainHistories Pairs of (domainRepoId, revisionId). This collection helps us find the * @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 * historical superordinate domain from its history entry and is obtained from calling {@link
@@ -530,59 +641,98 @@ public class RdePipeline implements Serializable {
private PCollection<KV<PendingDeposit, DepositFragment>> processSubordinateHosts( private PCollection<KV<PendingDeposit, DepositFragment>> processSubordinateHosts(
PCollection<KV<String, KV<String, CoGbkResult>>> superordinateDomains, PCollection<KV<String, KV<String, CoGbkResult>>> superordinateDomains,
PCollection<KV<String, Long>> domainHistories) { PCollection<KV<String, Long>> domainHistories) {
Counter subordinateHostFragmentCounter = Metrics.counter("RDE", "SubordinateHostFragment"); int batchSize = options.getHistoryEntryLoadBatchSize();
Counter referencedSubordinateHostCounter = Metrics.counter("RDE", "ReferencedSubordinateHost"); int numShards = options.getNumHistoryEntryShards();
return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT, superordinateDomains) return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID, superordinateDomains)
.and(REVISION_ID, domainHistories) .and(REVISION_ID, domainHistories)
.apply("Join Host:PendingDeposits with DomainHistory on Domain", CoGroupByKey.create()) .apply("Join Host:PendingDeposits with DomainHistory on Domain", CoGroupByKey.create())
.apply( .apply(
" Remove unreferenced Domain", "Remove Domains without subordinate hosts",
Filter.by( Filter.by(
kv -> { kv -> {
boolean toInclude = boolean toInclude =
kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT).iterator().hasNext() !Iterables.isEmpty(
&& kv.getValue().getAll(REVISION_ID).iterator().hasNext(); kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID))
&& !Iterables.isEmpty(kv.getValue().getAll(REVISION_ID));
if (toInclude) { if (toInclude) {
referencedSubordinateHostCounter.inc(); REFERENCED_SUBORDINATE_HOST_COUNTER.inc();
} }
return toInclude; 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.<Integer, KV<String, CoGbkResult>>of(
kv -> Math.floorMod(kv.getKey().hashCode(), numShards))
.withKeyType(integers()))
.apply(
"Group superordinate domains into batches",
GroupIntoBatches.<Integer, KV<String, CoGbkResult>>ofSize(batchSize).withShardedKey())
.apply( .apply(
"Map subordinate Host to DepositFragment", "Map subordinate Host to DepositFragment",
FlatMapElements.into( ParDo.of(
kvs( new DoFn<
TypeDescriptor.of(PendingDeposit.class), KV<ShardedKey<Integer>, Iterable<KV<String, CoGbkResult>>>,
TypeDescriptor.of(DepositFragment.class))) KV<PendingDeposit, DepositFragment>>() {
.via( @ProcessElement
(KV<String, CoGbkResult> kv) -> { public void processElement(
Domain superordinateDomain = @Element KV<ShardedKey<Integer>, Iterable<KV<String, CoGbkResult>>> element,
(Domain) OutputReceiver<KV<PendingDeposit, DepositFragment>> receiver) {
loadResourceByHistoryEntryId( ImmutableList<KV<String, CoGbkResult>> batchElements =
DomainHistory.class, ImmutableList.copyOf(element.getValue());
kv.getKey(), ImmutableSet<KV<String, Long>> domainKeys =
kv.getValue().getAll(REVISION_ID)); batchElements.stream()
ImmutableSet.Builder<KV<PendingDeposit, DepositFragment>> results = .map(
new ImmutableSet.Builder<>(); kv ->
KV.of(
kv.getKey(),
getSingleRevisionId(
DomainHistory.class,
kv.getKey(),
kv.getValue().getAll(REVISION_ID))))
.collect(toImmutableSet());
ImmutableSet<KV<String, Long>> 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<String, Domain> loadedDomains =
loadResourcesByHistoryEntryIds(
domainKeys, Domain.class, DomainHistory.class, watermark);
ImmutableMap<String, Host> loadedHosts =
loadResourcesByHistoryEntryIds(
hostKeys, Host.class, HostHistory.class, watermark);
for (KV<String, CoGbkResult> kv : batchElements) {
Domain superordinateDomain = loadedDomains.get(kv.getKey());
for (KV<String, CoGbkResult> hostToPendingDeposits : for (KV<String, CoGbkResult> hostToPendingDeposits :
kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT)) { kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID)) {
Host host = Host host = loadedHosts.get(hostToPendingDeposits.getKey());
(Host)
loadResourceByHistoryEntryId(
HostHistory.class,
hostToPendingDeposits.getKey(),
hostToPendingDeposits.getValue().getAll(REVISION_ID));
DepositFragment fragment = DepositFragment fragment =
marshaller.marshalSubordinateHost(host, superordinateDomain); marshaller.marshalSubordinateHost(host, superordinateDomain);
Streams.stream(hostToPendingDeposits.getValue().getAll(PENDING_DEPOSIT)) Streams.stream(hostToPendingDeposits.getValue().getAll(PENDING_DEPOSIT))
.distinct() .distinct()
.forEach( .forEach(
pendingDeposit -> { pendingDeposit -> {
subordinateHostFragmentCounter.inc(); SUBORDINATE_HOST_FRAGMENT_COUNTER.inc();
results.add(KV.of(pendingDeposit, fragment)); receiver.output(KV.of(pendingDeposit, fragment));
}); });
} }
return results.build(); }
})); }
}));
} }
/** /**
@@ -634,8 +784,8 @@ public class RdePipeline implements Serializable {
protected static final TupleTag<PendingDeposit> PENDING_DEPOSIT = new TupleTag<>() {}; protected static final TupleTag<PendingDeposit> PENDING_DEPOSIT = new TupleTag<>() {};
protected static final TupleTag<KV<String, CoGbkResult>> HOST_TO_PENDING_DEPOSIT = protected static final TupleTag<KV<String, CoGbkResult>>
new TupleTag<>() {}; HOST_TO_PENDING_DEPOSIT_AND_REVISION_ID = new TupleTag<>() {};
protected static final TupleTag<Long> REVISION_ID = new TupleTag<>() {}; protected static final TupleTag<Long> REVISION_ID = new TupleTag<>() {};
} }
core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java
+15 -1
View File
@@ -15,9 +15,10 @@
package google.registry.beam.rde; package google.registry.beam.rde;
import google.registry.beam.common.RegistryPipelineOptions; import google.registry.beam.common.RegistryPipelineOptions;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description; 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 { public interface RdePipelineOptions extends RegistryPipelineOptions {
@Description("The Base64-encoded serialized map of TLDs to PendingDeposit.") @Description("The Base64-encoded serialized map of TLDs to PendingDeposit.")
@@ -39,4 +40,17 @@ public interface RdePipelineOptions extends RegistryPipelineOptions {
String getStagingKey(); String getStagingKey();
void setStagingKey(String value); 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);
} }
core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json
+18
View File
@@ -43,6 +43,24 @@
"regexes": [ "regexes": [
"[A-Za-z0-9\\-_]+" "[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]*$"
]
} }
] ]
} }
core/src/test/java/google/registry/beam/rde/RdePipelineTest.java
+112 -29
View File
@@ -29,6 +29,7 @@ import static google.registry.rde.RdeResourceType.DOMAIN;
import static google.registry.rde.RdeResourceType.HOST; import static google.registry.rde.RdeResourceType.HOST;
import static google.registry.rde.RdeResourceType.REGISTRAR; import static google.registry.rde.RdeResourceType.REGISTRAR;
import static google.registry.testing.DatabaseHelper.createTld; 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.newDomain;
import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistActiveDomain;
import static google.registry.testing.DatabaseHelper.persistActiveHost; import static google.registry.testing.DatabaseHelper.persistActiveHost;
@@ -85,6 +86,7 @@ import google.registry.testing.FakeKeyringModule;
import java.io.IOException; import java.io.IOException;
import java.time.Duration; import java.time.Duration;
import java.time.Instant; import java.time.Instant;
import java.util.NoSuchElementException;
import java.util.function.Function; import java.util.function.Function;
import java.util.regex.Matcher; import java.util.regex.Matcher;
import java.util.regex.Pattern; import java.util.regex.Pattern;
@@ -166,36 +168,42 @@ public class RdePipelineTest {
.setReportAmount(1) .setReportAmount(1)
.build(); .build();
return persistResource( DomainHistory result =
new DomainHistory.Builder() persistResource(
.setType(HistoryEntry.Type.DOMAIN_CREATE) new DomainHistory.Builder()
.setXmlBytes("<xml></xml>".getBytes(UTF_8)) .setType(HistoryEntry.Type.DOMAIN_CREATE)
.setModificationTime(clock.now()) .setXmlBytes("<xml></xml>".getBytes(UTF_8))
.setRegistrarId("TheRegistrar") .setModificationTime(clock.now())
.setTrid(Trid.create("ABC-123", "server-trid")) .setRegistrarId("TheRegistrar")
.setBySuperuser(false) .setTrid(Trid.create("ABC-123", "server-trid"))
.setReason("reason") .setBySuperuser(false)
.setRequestedByRegistrar(true) .setReason("reason")
.setDomain(domain) .setRequestedByRegistrar(true)
.setDomainTransactionRecords(ImmutableSet.of(transactionRecord)) .setDomain(domain)
.setOtherRegistrarId("otherClient") .setDomainTransactionRecords(ImmutableSet.of(transactionRecord))
.setPeriod(Period.create(1, Period.Unit.YEARS)) .setOtherRegistrarId("otherClient")
.build()); .setPeriod(Period.create(1, Period.Unit.YEARS))
.build());
clock.advanceOneMilli();
return result;
} }
private HostHistory persistHostHistory(HostBase hostBase) { private HostHistory persistHostHistory(HostBase hostBase) {
return persistResource( HostHistory result =
new HostHistory.Builder() persistResource(
.setType(HistoryEntry.Type.HOST_CREATE) new HostHistory.Builder()
.setXmlBytes("<xml></xml>".getBytes(UTF_8)) .setType(HistoryEntry.Type.HOST_CREATE)
.setModificationTime(clock.now()) .setXmlBytes("<xml></xml>".getBytes(UTF_8))
.setRegistrarId("TheRegistrar") .setModificationTime(clock.now())
.setTrid(Trid.create("ABC-123", "server-trid")) .setRegistrarId("TheRegistrar")
.setBySuperuser(false) .setTrid(Trid.create("ABC-123", "server-trid"))
.setReason("reason") .setBySuperuser(false)
.setRequestedByRegistrar(true) .setReason("reason")
.setHost(hostBase) .setRequestedByRegistrar(true)
.build()); .setHost(hostBase)
.build());
clock.advanceOneMilli();
return result;
} }
@BeforeEach @BeforeEach
@@ -252,12 +260,12 @@ public class RdePipelineTest {
.build()); .build());
persistDomainHistory(kittyDomain); persistDomainHistory(kittyDomain);
// Should not appear because the TLD is not included in a pending deposit. // 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. // To be deleted.
Domain deletedDomain = persistActiveDomain("deleted.soy"); Domain deletedDomain = persistActiveDomain("deleted.soy");
persistDomainHistory(deletedDomain); persistDomainHistory(deletedDomain);
// Advance time // Advance time again just in case
clock.advanceOneMilli(); clock.advanceOneMilli();
persistDomainHistory(deletedDomain.asBuilder().setDeletionTime(clock.now()).build()); persistDomainHistory(deletedDomain.asBuilder().setDeletionTime(clock.now()).build());
kittyDomain = kittyDomain.asBuilder().setDomainName("cat.fun").build(); kittyDomain = kittyDomain.asBuilder().setDomainName("cat.fun").build();
@@ -425,6 +433,81 @@ public class RdePipelineTest {
pipeline.run().waitUntilFinish(); 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<String, Domain> 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<String, Domain> 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<String, Domain> 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 // The GCS folder listing can be a bit flaky, so retry if necessary
@RetryingTest(4) @RetryingTest(4)
void testSuccess_persistData() throws Exception { void testSuccess_persistData() throws Exception {