Simplify the CloudDnsWriter callWithRetry functional

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=176512218
This commit is contained in:
guyben 2017-11-21 08:33:09 -08:00 committed by jianglai
parent 3c43ece5be
commit 6f659659ff
2 changed files with 31 additions and 37 deletions

View file

@ -50,7 +50,6 @@ import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;
import javax.inject.Inject;
import javax.inject.Named;
import org.joda.time.Duration;
@ -272,47 +271,45 @@ public class CloudDnsWriter extends BaseDnsWriter {
*/
@Override
protected void commitUnchecked() {
retrier.callWithRetry(
getMutateZoneCallback(ImmutableMap.copyOf(desiredRecords)), ZoneStateException.class);
ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecordsCopy =
ImmutableMap.copyOf(desiredRecords);
retrier.callWithRetry(() -> mutateZone(desiredRecordsCopy), ZoneStateException.class);
logger.info("Wrote to Cloud DNS");
}
/**
* Get a callback to mutate the zone with the provided {@code desiredRecords}.
* Mutate the zone with the provided {@code desiredRecords}.
*/
@VisibleForTesting
Callable<Void> getMutateZoneCallback(
final ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecords) {
return () -> {
// Fetch all existing records for names that this writer is trying to modify
Builder<ResourceRecordSet> existingRecords = new Builder<>();
for (String domainName : desiredRecords.keySet()) {
List<ResourceRecordSet> existingRecordsForDomain = getResourceRecordsForDomain(domainName);
existingRecords.addAll(existingRecordsForDomain);
void mutateZone(ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecords)
throws IOException {
// Fetch all existing records for names that this writer is trying to modify
Builder<ResourceRecordSet> existingRecords = new Builder<>();
for (String domainName : desiredRecords.keySet()) {
List<ResourceRecordSet> existingRecordsForDomain = getResourceRecordsForDomain(domainName);
existingRecords.addAll(existingRecordsForDomain);
// Fetch glue records for in-bailiwick nameservers
for (ResourceRecordSet record : existingRecordsForDomain) {
if (!record.getType().equals("NS")) {
continue;
}
for (String hostName : record.getRrdatas()) {
if (hostName.endsWith(domainName) && !hostName.equals(domainName)) {
existingRecords.addAll(getResourceRecordsForDomain(hostName));
}
// Fetch glue records for in-bailiwick nameservers
for (ResourceRecordSet record : existingRecordsForDomain) {
if (!record.getType().equals("NS")) {
continue;
}
for (String hostName : record.getRrdatas()) {
if (hostName.endsWith(domainName) && !hostName.equals(domainName)) {
existingRecords.addAll(getResourceRecordsForDomain(hostName));
}
}
}
}
// Flatten the desired records into one set.
Builder<ResourceRecordSet> flattenedDesiredRecords = new Builder<>();
for (ImmutableSet<ResourceRecordSet> records : desiredRecords.values()) {
flattenedDesiredRecords.addAll(records);
}
// Flatten the desired records into one set.
Builder<ResourceRecordSet> flattenedDesiredRecords = new Builder<>();
for (ImmutableSet<ResourceRecordSet> records : desiredRecords.values()) {
flattenedDesiredRecords.addAll(records);
}
// Delete all existing records and add back the desired records
updateResourceRecords(flattenedDesiredRecords.build(), existingRecords.build());
return null;
};
// Delete all existing records and add back the desired records
updateResourceRecords(flattenedDesiredRecords.build(), existingRecords.build());
}
/**

View file

@ -22,6 +22,7 @@ import static google.registry.testing.DatastoreHelper.newDomainResource;
import static google.registry.testing.DatastoreHelper.newHostResource;
import static google.registry.testing.DatastoreHelper.persistResource;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@ -51,7 +52,6 @@ import java.io.IOException;
import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.util.concurrent.Callable;
import org.joda.time.Duration;
import org.junit.Before;
import org.junit.Rule;
@ -82,7 +82,6 @@ public class CloudDnsWriterTest {
@Mock private Dns.ResourceRecordSets.List listResourceRecordSetsRequest;
@Mock private Dns.Changes changes;
@Mock private Dns.Changes.Create createChangeRequest;
@Mock private Callable<Void> mutateZoneCallable;
@Captor ArgumentCaptor<String> recordNameCaptor;
@Captor ArgumentCaptor<String> zoneNameCaptor;
@Captor ArgumentCaptor<Change> changeCaptor;
@ -406,13 +405,11 @@ public class CloudDnsWriterTest {
@SuppressWarnings("unchecked")
public void retryMutateZoneOnError() throws Exception {
CloudDnsWriter spyWriter = spy(writer);
when(mutateZoneCallable.call()).thenThrow(ZoneStateException.class).thenReturn(null);
when(spyWriter.getMutateZoneCallback(
Matchers.any()))
.thenReturn(mutateZoneCallable);
// First call - throw. Second call - do nothing.
doThrow(ZoneStateException.class).doNothing().when(spyWriter).mutateZone(Matchers.any());
spyWriter.commit();
verify(mutateZoneCallable, times(2)).call();
verify(spyWriter, times(2)).mutateZone(Matchers.any());
}
@Test