> 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 {