From 24587491c96c89ad96156a70de6ac665911eccad Mon Sep 17 00:00:00 2001
From: mcilwain <mcilwain@google.com>
Date: Wed, 12 Jul 2017 08:11:51 -0700
Subject: [PATCH] Make re-save environment entities command use batching

This makes it take a lot less time to run (roughly a 10X speedup).

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=161666391
---
 .../ResaveEnvironmentEntitiesCommand.java     | 29 +++++++++++--------
 .../ResaveEnvironmentEntitiesCommandTest.java |  6 ++--
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/java/google/registry/tools/ResaveEnvironmentEntitiesCommand.java b/java/google/registry/tools/ResaveEnvironmentEntitiesCommand.java
index b207d4fe4..b2bc440d1 100644
--- a/java/google/registry/tools/ResaveEnvironmentEntitiesCommand.java
+++ b/java/google/registry/tools/ResaveEnvironmentEntitiesCommand.java
@@ -14,14 +14,13 @@
 
 package google.registry.tools;
 
-import static com.google.common.collect.Iterables.concat;
+import static com.google.common.collect.Lists.partition;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
 import com.beust.jcommander.Parameters;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.VoidWork;
-import google.registry.model.ImmutableObject;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registry.Registry;
@@ -36,19 +35,25 @@ import google.registry.tools.Command.RemoteApiCommand;
 @Parameters(commandDescription = "Re-save all environment entities.")
 final class ResaveEnvironmentEntitiesCommand implements RemoteApiCommand {
 
+  private static final int BATCH_SIZE = 10;
+
   @Override
   public void run() throws Exception {
-    Iterable<Key<? extends ImmutableObject>> keys = concat(
-        ofy().load().type(Registrar.class).ancestor(getCrossTldKey()).keys(),
-        ofy().load().type(Registry.class).ancestor(getCrossTldKey()).keys(),
-        ofy().load().type(RegistrarContact.class).ancestor(getCrossTldKey()).keys());
-    for (final Key<? extends ImmutableObject> key : keys) {
+    batchSave(Registry.class);
+    batchSave(Registrar.class);
+    batchSave(RegistrarContact.class);
+  }
+
+  private static <T> void batchSave(Class<T> clazz) {
+    System.out.printf("Re-saving %s entities.\n", clazz.getSimpleName());
+    for (final Iterable<Key<T>> batch :
+        partition(ofy().load().type(clazz).ancestor(getCrossTldKey()).keys().list(), BATCH_SIZE)) {
       ofy().transact(new VoidWork() {
-          @Override
-          public void vrun() {
-            ofy().save().entity(ofy().load().key(key).now());
-          }});
-      System.out.printf("Re-saved entity %s\n", key);
+        @Override
+        public void vrun() {
+          ofy().save().entities(ofy().load().keys(batch).values());
+        }});
+      System.out.printf("Re-saved entities batch: %s.\n", batch);
     }
   }
 }
diff --git a/javatests/google/registry/tools/ResaveEnvironmentEntitiesCommandTest.java b/javatests/google/registry/tools/ResaveEnvironmentEntitiesCommandTest.java
index e79e2ab9e..842242d81 100644
--- a/javatests/google/registry/tools/ResaveEnvironmentEntitiesCommandTest.java
+++ b/javatests/google/registry/tools/ResaveEnvironmentEntitiesCommandTest.java
@@ -55,9 +55,9 @@ public class ResaveEnvironmentEntitiesCommandTest
     assertThat(ofy().load().type(CommitLogMutation.class).keys()).isEmpty();
     runCommand();
 
-    // There are five entities that have been re-saved at this point (each in a separate
-    // transaction), so expect five manifests and five mutations.
-    assertThat(ofy().load().type(CommitLogManifest.class).keys()).hasSize(5);
+    // There are 5 entities that have been re-saved at this point (in 3 transactions, one for each
+    // type), so expect 3 manifests and 5 mutations.
+    assertThat(ofy().load().type(CommitLogManifest.class).keys()).hasSize(3);
     Iterable<ImmutableObject> savedEntities =
         transform(
             ofy().load().type(CommitLogMutation.class).list(),