From 2995c7f8832fc560b0103ea499d8358fbd5d57ef Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Wed, 29 Jul 2026 17:01:45 -0400 Subject: [PATCH] Batch database queries in RdePipeline 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. In loadResourcesByHistoryEntryIds we can query directly by the revisionIds because those are guaranteed to be unique. Self-scan D.2 number 5 --- .../google/registry/beam/rde/RdePipeline.java | 387 +++++++++++------- .../registry/beam/rde/RdePipelineOptions.java | 10 +- .../registry/beam/rde/RdePipelineTest.java | 88 ++-- 3 files changed, 309 insertions(+), 176 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 37507988e1b..d76714b8e17 100644 --- a/core/src/main/java/google/registry/beam/rde/RdePipeline.java +++ b/core/src/main/java/google/registry/beam/rde/RdePipeline.java @@ -16,6 +16,8 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static google.registry.beam.rde.RdePipeline.TupleTags.DOMAIN_FRAGMENTS; import static google.registry.beam.rde.RdePipeline.TupleTags.EXTERNAL_HOST_FRAGMENTS; @@ -24,12 +26,12 @@ import static google.registry.beam.rde.RdePipeline.TupleTags.REFERENCED_HOSTS; import static google.registry.beam.rde.RdePipeline.TupleTags.REVISION_ID; import static google.registry.beam.rde.RdePipeline.TupleTags.SUPERORDINATE_DOMAINS; -import static google.registry.model.reporting.HistoryEntryDao.RESOURCE_TYPES_TO_HISTORY_TYPES; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.util.SafeSerializationUtils.safeDeserializeCollection; import static google.registry.util.SafeSerializationUtils.serializeCollection; import static google.registry.util.SerializeUtils.decodeBase64; import static google.registry.util.SerializeUtils.encodeBase64; +import static org.apache.beam.sdk.values.TypeDescriptors.integers; import static org.apache.beam.sdk.values.TypeDescriptors.kvs; import com.google.common.collect.ImmutableList; @@ -57,7 +59,6 @@ import google.registry.model.registrar.Registrar; import google.registry.model.registrar.Registrar.Type; import google.registry.model.reporting.HistoryEntry; -import google.registry.model.reporting.HistoryEntry.HistoryEntryId; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; import google.registry.persistence.VKey; import google.registry.rde.DepositFragment; @@ -71,6 +72,7 @@ import java.io.IOException; import java.io.Serializable; import java.time.Instant; +import java.util.NoSuchElementException; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.KvCoder; @@ -85,10 +87,13 @@ import org.apache.beam.sdk.transforms.FlatMapElements; import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.GroupIntoBatches; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.join.CoGbkResult; import org.apache.beam.sdk.transforms.join.CoGroupByKey; import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; @@ -327,18 +332,17 @@ private PCollection> getMostRecentHist .withCoder(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))); } - private EppResource loadResourceByHistoryEntryId( - Class historyEntryClazz, String repoId, Iterable revisionIds) { + private static long getSingleRevisionId( + Class historyEntryClazz, String repoId, Iterable revisionIds) { ImmutableList ids = ImmutableList.copyOf(revisionIds); - // The size should always be 1 because we are only getting one repo ID -> revision ID pair per - // repo ID from the source transform (the JPA query in the method above). But for some reason - // 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. if (ids.size() != 1) { ImmutableSet dedupedIds = ImmutableSet.copyOf(ids); checkState( @@ -351,22 +355,52 @@ private EppResource loadResourceByHistoryEntryId( "Duplicate revision IDs detected for %s repo ID %s: %s", historyEntryClazz.getSimpleName(), repoId, ids); } - return loadResourceByHistoryEntryId(historyEntryClazz, repoId, ids.get(0)); + return ids.getFirst(); } - private EppResource loadResourceByHistoryEntryId( - Class historyEntryClazz, String repoId, long revisionId) { - return tm().transact( - () -> - tm().loadByKey( - VKey.create(historyEntryClazz, new HistoryEntryId(repoId, revisionId)))) - .getResourceAtPointInTime() - .map(resource -> resource.cloneProjectedAtTime(watermark)) - .get(); + + ImmutableMap loadResourcesByHistoryEntryIds( + Iterable> repoAndRevisionIds, + Class resourceClass, + Class historyEntryClass) { + ImmutableList> ids = ImmutableList.copyOf(repoAndRevisionIds); + if (ids.isEmpty()) { + return ImmutableMap.of(); + } + ImmutableSet expectedRepoIds = ids.stream().map(KV::getKey).collect(toImmutableSet()); + ImmutableList revisionIds = ids.stream().map(KV::getValue).collect(toImmutableList()); + ImmutableMap result = + tm().transact( + () -> + tm().query( + String.format( + "FROM %s WHERE revisionId IN (:revisionIds)", + historyEntryClass.getSimpleName()), + historyEntryClass) + .setParameter("revisionIds", revisionIds) + .getResultStream() + .collect( + toImmutableMap( + HistoryEntry::getRepoId, + entry -> + entry + .getResourceAtPointInTime() + .map(r -> r.cloneProjectedAtTime(watermark)) + .map(resourceClass::cast) + .get()))); + // Fail fast on items being missing unexpectedly + if (result.size() != expectedRepoIds.size()) { + 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 @@ -378,38 +412,30 @@ private EppResource loadResourceByHistoryEntryId( * @return a pair of (repoId, ([pendingDeposit], [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) - .apply( - String.format( - "Join PendingDeposit with %s revision ID on %s", historyEntryName, resourceName), - CoGroupByKey.create()) + private static PCollection> removeUnreferencedHost( + PCollection> referencedHosts, + PCollection> hostHistories) { + Counter referencedHostCounter = Metrics.counter("RDE", "Referenced Host"); + return KeyedPCollectionTuple.of(PENDING_DEPOSIT, referencedHosts) + .and(REVISION_ID, hostHistories) + .apply("Join PendingDeposit with HostHistory revision ID on Host", CoGroupByKey.create()) .apply( - String.format("Remove unreferenced %s", resourceName), + "Remove unreferenced Host", Filter.by( (KV kv) -> { boolean toInclude = - // If a resource does not have corresponding pending deposit, it is not - // referenced and should not be included. + // If a host 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 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(); if (toInclude) { - referencedResourceCounter.inc(); + referencedHostCounter.inc(); } return toInclude; })); @@ -419,50 +445,65 @@ private PCollectionTuple processDomainHistories(PCollection> do Counter activeDomainCounter = Metrics.counter("RDE", "ActiveDomainBase"); Counter domainFragmentCounter = Metrics.counter("RDE", "DomainFragment"); Counter referencedHostCounter = Metrics.counter("RDE", "ReferencedHost"); - return domainHistories.apply( - "Map DomainHistory to DepositFragment and emit referenced Host", - ParDo.of( - new DoFn, KV>() { - @ProcessElement - public void processElement( - @Element KV kv, MultiOutputReceiver receiver) { - activeDomainCounter.inc(); - Domain domain = - (Domain) - loadResourceByHistoryEntryId( - DomainHistory.class, kv.getKey(), kv.getValue()); - pendingDeposits.stream() - .filter(pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) - .forEach( - pendingDeposit -> { - // Domains are always deposited in both modes. - domainFragmentCounter.inc(); - receiver - .get(DOMAIN_FRAGMENTS) - .output( - KV.of( - pendingDeposit, - marshaller.marshalDomain(domain, pendingDeposit.mode()))); - // Hosts are only deposited in RDE, not BRDA. - if (pendingDeposit.mode() == RdeMode.FULL) { - if (domain.getNsHosts() != null) { - referencedHostCounter.inc(domain.getNsHosts().size()); - domain - .getNsHosts() - .forEach( - hostKey -> - receiver - .get(REFERENCED_HOSTS) - .output( - KV.of( - (String) hostKey.getKey(), - pendingDeposit))); - } - } - }); - } - }) - .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); + return domainHistories + .apply( + "Assign all domain histories a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group domain histories into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) + .apply( + "Map DomainHistory to DepositFragment and emit referenced Host", + ParDo.of( + new DoFn< + KV, Iterable>>, + KV>() { + @ProcessElement + public void processElement( + @Element KV, Iterable>> element, + MultiOutputReceiver receiver) { + ImmutableMap loadedDomains = + loadResourcesByHistoryEntryIds( + element.getValue(), Domain.class, DomainHistory.class); + for (KV kv : element.getValue()) { + Domain domain = loadedDomains.get(kv.getKey()); + activeDomainCounter.inc(); + pendingDeposits.stream() + .filter( + pendingDeposit -> pendingDeposit.tld().equals(domain.getTld())) + .forEach( + pendingDeposit -> { + // Domains are always deposited in both modes. + domainFragmentCounter.inc(); + receiver + .get(DOMAIN_FRAGMENTS) + .output( + KV.of( + pendingDeposit, + marshaller.marshalDomain( + domain, pendingDeposit.mode()))); + // Hosts are only deposited in RDE, not BRDA. + if (pendingDeposit.mode() == RdeMode.FULL) { + if (domain.getNsHosts() != null) { + referencedHostCounter.inc(domain.getNsHosts().size()); + domain + .getNsHosts() + .forEach( + hostKey -> + receiver + .get(REFERENCED_HOSTS) + .output( + KV.of( + (String) hostKey.getKey(), + pendingDeposit))); + } + } + }); + } + } + }) + .withOutputTags(DOMAIN_FRAGMENTS, TupleTagList.of(REFERENCED_HOSTS))); } private PCollectionTuple processHostHistories( @@ -471,45 +512,66 @@ private PCollectionTuple processHostHistories( Counter subordinateHostCounter = Metrics.counter("RDE", "SubordinateHost"); Counter externalHostCounter = Metrics.counter("RDE", "ExternalHost"); Counter externalHostFragmentCounter = Metrics.counter("RDE", "ExternalHostFragment"); - return removeUnreferencedResource(referencedHosts, hostHistories, Host.class) + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); + return removeUnreferencedHost(referencedHosts, hostHistories) + .apply( + "Assign all host histories a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group referenced hosts into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) .apply( "Map external DomainResource to DepositFragment and process subordinate domains", ParDo.of( - new DoFn, KV>() { + new DoFn< + KV, Iterable>>, + KV>() { @ProcessElement public void processElement( - @Element KV kv, MultiOutputReceiver receiver) { - Host host = - (Host) - loadResourceByHistoryEntryId( - HostHistory.class, - kv.getKey(), - kv.getValue().getAll(REVISION_ID)); - // When a host is subordinate, we need to find its superordinate domain and - // include it in the deposit as well. - if (host.isSubordinate()) { - subordinateHostCounter.inc(); - receiver - .get(SUPERORDINATE_DOMAINS) - .output( - // The output are pairs of - // (superordinateDomainRepoId, - // (subordinateHostRepoId, (pendingDeposit, revisionId))). - KV.of((String) host.getSuperordinateDomain().getKey(), kv)); - } else { - externalHostCounter.inc(); - DepositFragment fragment = marshaller.marshalExternalHost(host); - Streams.stream(kv.getValue().getAll(PENDING_DEPOSIT)) - // The same host could be used by multiple domains, therefore - // matched to the same pending deposit multiple times. - .distinct() - .forEach( - pendingDeposit -> { - externalHostFragmentCounter.inc(); - receiver - .get(EXTERNAL_HOST_FRAGMENTS) - .output(KV.of(pendingDeposit, fragment)); - }); + @Element + KV, Iterable>> element, + MultiOutputReceiver receiver) { + ImmutableSet> hostKeys = + Streams.stream(element.getValue()) + .map( + kv -> + KV.of( + kv.getKey(), + getSingleRevisionId( + HostHistory.class, + kv.getKey(), + kv.getValue().getAll(REVISION_ID)))) + .collect(toImmutableSet()); + ImmutableMap loadedHosts = + loadResourcesByHistoryEntryIds(hostKeys, Host.class, HostHistory.class); + for (KV kv : element.getValue()) { + Host host = loadedHosts.get(kv.getKey()); + // When a host is subordinate, we need to find its superordinate domain + // and include it in the deposit as well. + if (host.isSubordinate()) { + subordinateHostCounter.inc(); + receiver + .get(SUPERORDINATE_DOMAINS) + .output( + // The output are pairs of + // (superordinateDomainRepoId, + // (subordinateHostRepoId, (pendingDeposit, revisionId))). + KV.of((String) host.getSuperordinateDomain().getKey(), kv)); + } else { + externalHostCounter.inc(); + DepositFragment fragment = marshaller.marshalExternalHost(host); + Streams.stream(kv.getValue().getAll(PENDING_DEPOSIT)) + // The same host could be used by multiple domains, therefore + // matched to the same pending deposit multiple times. + .distinct() + .forEach( + pendingDeposit -> { + externalHostFragmentCounter.inc(); + receiver + .get(EXTERNAL_HOST_FRAGMENTS) + .output(KV.of(pendingDeposit, fragment)); + }); + } } } }) @@ -532,6 +594,7 @@ private PCollection> processSubordinateHosts PCollection> domainHistories) { Counter subordinateHostFragmentCounter = Metrics.counter("RDE", "SubordinateHostFragment"); Counter referencedSubordinateHostCounter = Metrics.counter("RDE", "ReferencedSubordinateHost"); + int batchSize = options.getRdeHistoryEntryLoadBatchSize(); return KeyedPCollectionTuple.of(HOST_TO_PENDING_DEPOSIT, superordinateDomains) .and(REVISION_ID, domainHistories) .apply("Join Host:PendingDeposits with DomainHistory on Domain", CoGroupByKey.create()) @@ -547,30 +610,59 @@ private PCollection> processSubordinateHosts } return toInclude; })) + .apply( + "Assign all subordinate hosts a key of 0 so they can be batched", + WithKeys.>of(0).withKeyType(integers())) + .apply( + "Group subordinate hosts into batches", + GroupIntoBatches.>ofSize(batchSize).withShardedKey()) .apply( "Map subordinate Host to DepositFragment", - FlatMapElements.into( - kvs( - TypeDescriptor.of(PendingDeposit.class), - TypeDescriptor.of(DepositFragment.class))) - .via( - (KV kv) -> { - Domain superordinateDomain = - (Domain) - loadResourceByHistoryEntryId( - DomainHistory.class, - kv.getKey(), - kv.getValue().getAll(REVISION_ID)); - ImmutableSet.Builder> results = - new ImmutableSet.Builder<>(); + ParDo.of( + new DoFn< + KV, Iterable>>, + KV>() { + @ProcessElement + public void processElement( + @Element KV, Iterable>> element, + OutputReceiver> receiver) { + ImmutableSet> domainKeys = + Streams.stream(element.getValue()) + .map( + kv -> + KV.of( + kv.getKey(), + getSingleRevisionId( + DomainHistory.class, + kv.getKey(), + kv.getValue().getAll(REVISION_ID)))) + .collect(toImmutableSet()); + ImmutableSet> hostKeys = + Streams.stream(element.getValue()) + .flatMap( + kv -> + Streams.stream(kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT)) + .map( + hostToPendingDeposits -> + KV.of( + hostToPendingDeposits.getKey(), + getSingleRevisionId( + HostHistory.class, + hostToPendingDeposits.getKey(), + hostToPendingDeposits + .getValue() + .getAll(REVISION_ID))))) + .collect(toImmutableSet()); + ImmutableMap loadedDomains = + loadResourcesByHistoryEntryIds( + domainKeys, Domain.class, DomainHistory.class); + ImmutableMap loadedHosts = + loadResourcesByHistoryEntryIds(hostKeys, Host.class, HostHistory.class); + for (KV kv : element.getValue()) { + Domain superordinateDomain = loadedDomains.get(kv.getKey()); for (KV hostToPendingDeposits : kv.getValue().getAll(HOST_TO_PENDING_DEPOSIT)) { - Host host = - (Host) - loadResourceByHistoryEntryId( - HostHistory.class, - hostToPendingDeposits.getKey(), - hostToPendingDeposits.getValue().getAll(REVISION_ID)); + Host host = loadedHosts.get(hostToPendingDeposits.getKey()); DepositFragment fragment = marshaller.marshalSubordinateHost(host, superordinateDomain); Streams.stream(hostToPendingDeposits.getValue().getAll(PENDING_DEPOSIT)) @@ -578,11 +670,12 @@ private PCollection> processSubordinateHosts .forEach( pendingDeposit -> { subordinateHostFragmentCounter.inc(); - results.add(KV.of(pendingDeposit, fragment)); + receiver.output(KV.of(pendingDeposit, fragment)); }); } - return results.build(); - })); + } + } + })); } /** diff --git a/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java b/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java index c365494b67b..e3601e198d2 100644 --- a/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java +++ b/core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java @@ -15,9 +15,10 @@ package google.registry.beam.rde; import google.registry.beam.common.RegistryPipelineOptions; +import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; -/** Custom options for running the spec11 pipeline. */ +/** Custom options for running the RDE pipeline. */ public interface RdePipelineOptions extends RegistryPipelineOptions { @Description("The Base64-encoded serialized map of TLDs to PendingDeposit.") @@ -39,4 +40,11 @@ public interface RdePipelineOptions extends RegistryPipelineOptions { String getStagingKey(); void setStagingKey(String value); + + @Description( + "The number of history entries to batch load from the SQL database in one operation.") + @Default.Integer(500) + int getRdeHistoryEntryLoadBatchSize(); + + void setRdeHistoryEntryLoadBatchSize(int value); } diff --git a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java index bd8fa1314fc..de2ff0c6a5a 100644 --- a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java +++ b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java @@ -85,6 +85,7 @@ import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.NoSuchElementException; import java.util.function.Function; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -166,36 +167,42 @@ private DomainHistory persistDomainHistory(DomainBase domain) { .setReportAmount(1) .build(); - return persistResource( - new DomainHistory.Builder() - .setType(HistoryEntry.Type.DOMAIN_CREATE) - .setXmlBytes("".getBytes(UTF_8)) - .setModificationTime(clock.now()) - .setRegistrarId("TheRegistrar") - .setTrid(Trid.create("ABC-123", "server-trid")) - .setBySuperuser(false) - .setReason("reason") - .setRequestedByRegistrar(true) - .setDomain(domain) - .setDomainTransactionRecords(ImmutableSet.of(transactionRecord)) - .setOtherRegistrarId("otherClient") - .setPeriod(Period.create(1, Period.Unit.YEARS)) - .build()); + DomainHistory result = + persistResource( + new DomainHistory.Builder() + .setType(HistoryEntry.Type.DOMAIN_CREATE) + .setXmlBytes("".getBytes(UTF_8)) + .setModificationTime(clock.now()) + .setRegistrarId("TheRegistrar") + .setTrid(Trid.create("ABC-123", "server-trid")) + .setBySuperuser(false) + .setReason("reason") + .setRequestedByRegistrar(true) + .setDomain(domain) + .setDomainTransactionRecords(ImmutableSet.of(transactionRecord)) + .setOtherRegistrarId("otherClient") + .setPeriod(Period.create(1, Period.Unit.YEARS)) + .build()); + clock.advanceOneMilli(); + return result; } private HostHistory persistHostHistory(HostBase hostBase) { - return persistResource( - new HostHistory.Builder() - .setType(HistoryEntry.Type.HOST_CREATE) - .setXmlBytes("".getBytes(UTF_8)) - .setModificationTime(clock.now()) - .setRegistrarId("TheRegistrar") - .setTrid(Trid.create("ABC-123", "server-trid")) - .setBySuperuser(false) - .setReason("reason") - .setRequestedByRegistrar(true) - .setHost(hostBase) - .build()); + HostHistory result = + persistResource( + new HostHistory.Builder() + .setType(HistoryEntry.Type.HOST_CREATE) + .setXmlBytes("".getBytes(UTF_8)) + .setModificationTime(clock.now()) + .setRegistrarId("TheRegistrar") + .setTrid(Trid.create("ABC-123", "server-trid")) + .setBySuperuser(false) + .setReason("reason") + .setRequestedByRegistrar(true) + .setHost(hostBase) + .build()); + clock.advanceOneMilli(); + return result; } @BeforeEach @@ -257,7 +264,7 @@ void beforeEach() throws Exception { Domain deletedDomain = persistActiveDomain("deleted.soy"); persistDomainHistory(deletedDomain); - // Advance time + // Advance time again just in case clock.advanceOneMilli(); persistDomainHistory(deletedDomain.asBuilder().setDeletionTime(clock.now()).build()); kittyDomain = kittyDomain.asBuilder().setDomainName("cat.fun").build(); @@ -425,6 +432,31 @@ void testSuccess_createFragments() { pipeline.run().waitUntilFinish(); } + @Test + void testSuccess_createFragments_smallBatchSize() { + options.setRdeHistoryEntryLoadBatchSize(1); + testSuccess_createFragments(); + } + + @Test + void testSuccess_createFragments_multiBatch() { + options.setRdeHistoryEntryLoadBatchSize(2); + testSuccess_createFragments(); + } + + @Test + void testFailure_missingHistoryEntry() { + NoSuchElementException thrown = + assertThrows( + NoSuchElementException.class, + () -> + rdePipeline.loadResourcesByHistoryEntryIds( + ImmutableList.of(KV.of("nonexistent", 12345L)), + Domain.class, + DomainHistory.class)); + assertThat(thrown).hasMessageThat().contains("nonexistent"); + } + // The GCS folder listing can be a bit flaky, so retry if necessary @RetryingTest(4) void testSuccess_persistData() throws Exception {