Add retry logic to CreateLrpTokensCommand

A transient 404 on entity save interrupts a long-running run of this command.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=138400654

java/google/registry/tools/BUILD

@@ -66,6 +66,7 @@ java_library(
         "//java/google/registry/xjc",
         "//java/google/registry/xml",
         "//third_party/java/appengine:appengine-api",
+        "//third_party/java/appengine:appengine-remote-api-link",
         "//third_party/java/bouncycastle",
         "//third_party/java/bouncycastle_bcpg",
         "//third_party/java/dagger",

java/google/registry/tools/CreateLrpTokensCommand.java

@@ -24,6 +24,8 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.google.appengine.tools.remoteapi.RemoteApiException;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.CharMatcher;
 import com.google.common.base.Function;
 import com.google.common.base.Splitter;
@@ -40,24 +42,28 @@ import google.registry.tools.Command.RemoteApiCommand;
 import google.registry.tools.params.KeyValueMapParameter.StringToIntegerMap;
 import google.registry.tools.params.KeyValueMapParameter.StringToStringMap;
 import google.registry.tools.params.PathParameter;
+import google.registry.util.NonFinalForTesting;
+import google.registry.util.Retrier;
 import google.registry.util.StringGenerator;
 import google.registry.util.TokenUtils;
 import java.io.StringReader;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.Set;
+import java.util.concurrent.Callable;
 import javax.inject.Inject;
 
 /**
  * Command to create one or more LRP tokens, given assignee(s) as either a parameter or a text file.
  */
+@NonFinalForTesting
 @Parameters(
     separators = " =",
     commandDescription = "Create an LRP token for a given assignee (using -a) or import a text"
         + " file of assignees for bulk token creation (using -i). Assignee/token pairs are printed"
         + " to stdout, and should be piped to a file for distribution to assignees or for cleanup"
         + " in the event of a command interruption.")
-public final class CreateLrpTokensCommand implements RemoteApiCommand {
+public class CreateLrpTokensCommand implements RemoteApiCommand {
 
   @Parameter(
       names = {"-a", "--assignee"},
@@ -95,6 +101,7 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand {
   private ImmutableMap<String, Integer> metadataColumns;
 
   @Inject StringGenerator stringGenerator;
+  @Inject Retrier retrier;
 
   private static final int BATCH_SIZE = 20;
   
@@ -126,7 +133,7 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand {
 
     String line = null;
     do {
-      ImmutableSet.Builder<LrpTokenEntity> tokensToSave = new ImmutableSet.Builder<>();
+      ImmutableSet.Builder<LrpTokenEntity> tokensToSaveBuilder = new ImmutableSet.Builder<>();
       for (String token : generateTokens(BATCH_SIZE)) {
         line = reader.readLine();
         if (!isNullOrEmpty(line)) {
@@ -155,14 +162,22 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand {
             }
             tokenBuilder.setMetadata(metadataBuilder.build());
           }
-          tokensToSave.add(tokenBuilder.build());
+          tokensToSaveBuilder.add(tokenBuilder.build());
         }
       }
-      saveTokens(tokensToSave.build());
+      final ImmutableSet<LrpTokenEntity> tokensToSave = tokensToSaveBuilder.build();
+      // Wrap in a retrier to deal with transient 404 errors (thrown as RemoteApiExceptions).
+      retrier.callWithRetry(new Callable<Void>() {
+        @Override
+        public Void call() throws Exception {
+          saveTokens(tokensToSave);
+          return null;
+        }}, RemoteApiException.class);
     } while (line != null);
   }
 
-  private void saveTokens(final ImmutableSet<LrpTokenEntity> tokens) {
+  @VisibleForTesting
+  void saveTokens(final ImmutableSet<LrpTokenEntity> tokens) {
     Collection<LrpTokenEntity> savedTokens =
         ofy().transact(new Work<Collection<LrpTokenEntity>>() {
           @Override

java/google/registry/tools/RegistryToolComponent.java

@@ -25,6 +25,7 @@ import google.registry.request.Modules.DatastoreServiceModule;
 import google.registry.request.Modules.Jackson2Module;
 import google.registry.request.Modules.URLFetchServiceModule;
 import google.registry.util.SystemClock.SystemClockModule;
+import google.registry.util.SystemSleeper.SystemSleeperModule;
 
 /**
  * Dagger component for Registry Tool.
@@ -43,6 +44,7 @@ import google.registry.util.SystemClock.SystemClockModule;
     KeyModule.class,
     RegistryToolModule.class,
     SystemClockModule.class,
+    SystemSleeperModule.class,
     URLFetchServiceModule.class,
     VoidDnsWriterModule.class,
   }

java/google/registry/util/NonFinalForTesting.java

@@ -16,6 +16,7 @@ package google.registry.util;
 
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.TYPE;
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import java.lang.annotation.Documented;
@@ -36,5 +37,5 @@ import java.lang.annotation.Target;
  */
 @Documented
 @Retention(SOURCE)
-@Target({FIELD, METHOD})
+@Target({FIELD, METHOD, TYPE})
 public @interface NonFinalForTesting {}

javatests/google/registry/tools/BUILD

@@ -24,6 +24,7 @@ java_library(
         "//java/com/google/common/net",
         "//java/com/google/common/reflect",
         "//third_party/java/appengine:appengine-api-testonly",
+        "//third_party/java/appengine:appengine-remote-api-link",
         "//third_party/java/jcommander",
         "//third_party/java/joda_money",
         "//third_party/java/joda_time",

javatests/google/registry/tools/CreateLrpTokensCommandTest.java

@@ -19,7 +19,11 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
 
+import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.Files;
@@ -28,10 +32,14 @@ import google.registry.model.domain.LrpTokenEntity;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.DeterministicStringGenerator;
 import google.registry.testing.DeterministicStringGenerator.Rule;
+import google.registry.testing.FakeClock;
+import google.registry.testing.FakeSleeper;
+import google.registry.util.Retrier;
 import java.io.File;
 import java.io.IOException;
 import java.util.Set;
 import javax.annotation.Nullable;
+import org.joda.time.DateTime;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -48,6 +56,8 @@ public class CreateLrpTokensCommandTest extends CommandTestCase<CreateLrpTokensC
     assigneeFile = tmpDir.newFile("lrp_assignees.txt");
     assigneeFilePath = assigneeFile.getPath();
     command.stringGenerator = stringGenerator;
+    command.retrier =
+        new Retrier(new FakeSleeper(new FakeClock(DateTime.parse("2000-01-01TZ"))), 3);
     createTld("tld");
   }
 
@@ -59,6 +69,21 @@ public class CreateLrpTokensCommandTest extends CommandTestCase<CreateLrpTokensC
     assertInStdout("domain.tld,LRP_abcdefghijklmnop");
   }
 
+  @Test
+  public void testSuccess_oneAssignee_retry() throws Exception {
+    CreateLrpTokensCommand spyCommand = spy(command);
+    RemoteApiException fakeException = new RemoteApiException("foo", "foo", "foo", new Exception());
+    doThrow(fakeException)
+        .doThrow(fakeException)
+        .doCallRealMethod()
+        .when(spyCommand)
+        .saveTokens(any());
+    runCommand("--assignee=domain.tld", "--tlds=tld");
+    assertLrpTokens(
+        createToken("LRP_abcdefghijklmnop", "domain.tld", ImmutableSet.of("tld"), null, null));
+    assertInStdout("domain.tld,LRP_abcdefghijklmnop");
+  }
+
   @Test
   public void testSuccess_oneAssignee_withMetadata() throws Exception {
     runCommand("--assignee=domain.tld", "--tlds=tld", "--metadata=key=foo,key2=bar");