Expand AckPollMessagesCommand to ack PollMessage.Autorenew (#647)

* Expand AckPollMessagesCommand to ack PollMessage.Autorenew

* Rebase on master and address comment

* Resolve comments

core/src/main/java/google/registry/flows/poll/PollAckFlow.java

@@ -14,8 +14,8 @@
 
 package google.registry.flows.poll;
 
-import static com.google.common.base.Preconditions.checkState;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
+import static google.registry.flows.poll.PollFlowUtils.ackPollMessage;
 import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -100,27 +100,8 @@ public class PollAckFlow implements TransactionalFlow {
     // This keeps track of whether we should include the current acked message in the updated
     // message count that's returned to the user. The only case where we do so is if an autorenew
     // poll message is acked, but its next event is already ready to be delivered.
-    boolean includeAckedMessageInCount = false;
+    boolean includeAckedMessageInCount = ackPollMessage(pollMessage);
-    if (pollMessage instanceof PollMessage.OneTime) {
-      // One-time poll messages are deleted once acked.
-      ofy().delete().entity(pollMessage);
-    } else {
-      checkState(pollMessage instanceof PollMessage.Autorenew, "Unknown poll message type");
-      PollMessage.Autorenew autorenewPollMessage = (PollMessage.Autorenew) pollMessage;
 
-      // Move the eventTime of this autorenew poll message forward by a year.
-      DateTime nextEventTime = autorenewPollMessage.getEventTime().plusYears(1);
-
-      // If the next event falls within the bounds of the end time, then just update the eventTime
-      // and re-save it for future autorenew poll messages to be delivered. Otherwise, this
-      // autorenew poll message has no more events to deliver and should be deleted.
-      if (nextEventTime.isBefore(autorenewPollMessage.getAutorenewEndTime())) {
-        ofy().save().entity(autorenewPollMessage.asBuilder().setEventTime(nextEventTime).build());
-        includeAckedMessageInCount = isBeforeOrAt(nextEventTime, now);
-      } else {
-        ofy().delete().entity(autorenewPollMessage);
-      }
-    }
     // We need to return the new queue length. If this was the last message in the queue being
     // acked, then we return a special status code indicating that. Note that the query will
     // include the message being acked.

core/src/main/java/google/registry/flows/poll/PollFlowUtils.java

@@ -14,7 +14,10 @@
 
 package google.registry.flows.poll;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.poll.PollMessage;
@@ -33,4 +36,42 @@ public final class PollFlowUtils {
         .filter("eventTime <=", now.toDate())
         .order("eventTime");
   }
+
+  /**
+   * Acknowledges the given {@link PollMessage} and returns whether we should include the current
+   * acked message in the updated message count that's returned to the user.
+   *
+   * <p>The only case where we do so is if an autorenew poll message is acked, but its next event is
+   * already ready to be delivered.
+   */
+  public static boolean ackPollMessage(PollMessage pollMessage) {
+    checkArgument(
+        isBeforeOrAt(pollMessage.getEventTime(), tm().getTransactionTime()),
+        "Cannot ACK poll message with ID %s because its event time is in the future: %s",
+        pollMessage.getId(),
+        pollMessage.getEventTime());
+    boolean includeAckedMessageInCount = false;
+    if (pollMessage instanceof PollMessage.OneTime) {
+      // One-time poll messages are deleted once acked.
+      tm().delete(pollMessage.createVKey());
+    } else if (pollMessage instanceof PollMessage.Autorenew) {
+      PollMessage.Autorenew autorenewPollMessage = (PollMessage.Autorenew) pollMessage;
+
+      // Move the eventTime of this autorenew poll message forward by a year.
+      DateTime nextEventTime = autorenewPollMessage.getEventTime().plusYears(1);
+
+      // If the next event falls within the bounds of the end time, then just update the eventTime
+      // and re-save it for future autorenew poll messages to be delivered. Otherwise, this
+      // autorenew poll message has no more events to deliver and should be deleted.
+      if (nextEventTime.isBefore(autorenewPollMessage.getAutorenewEndTime())) {
+        tm().saveNewOrUpdate(autorenewPollMessage.asBuilder().setEventTime(nextEventTime).build());
+        includeAckedMessageInCount = isBeforeOrAt(nextEventTime, tm().getTransactionTime());
+      } else {
+        tm().delete(autorenewPollMessage.createVKey());
+      }
+    } else {
+      throw new IllegalArgumentException("Unknown poll message type: " + pollMessage.getClass());
+    }
+    return includeAckedMessageInCount;
+  }
 }

core/src/main/java/google/registry/tools/AckPollMessagesCommand.java

@@ -28,6 +28,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.cmd.QueryKeys;
+import google.registry.flows.poll.PollFlowUtils;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessage.Autorenew;
 import google.registry.model.poll.PollMessage.OneTime;
@@ -49,8 +50,7 @@ import javax.inject.Inject;
  * delete confirmations) and it is desired that the registrar be able to ACK the rest of the poll
  * messages in-band.
  *
- * <p>This command only ACKs {@link OneTime} poll messages because that's our only use case so far,
+ * <p>This command ACKs both {@link OneTime} and {@link Autorenew} poll messages. The main
- * but it could be extended to ACK {@link Autorenew} poll messages as well if needed. The main
  * difference is that one-time poll messages are deleted when ACKed whereas Autorenews are sometimes
  * modified and re-saved instead, if the corresponding domain is still active.
  *
@@ -85,18 +85,17 @@ final class AckPollMessagesCommand implements CommandWithRemoteApi {
   @Override
   public void run() {
     QueryKeys<PollMessage> query = getPollMessagesQuery(clientId, clock.nowUtc()).keys();
+    // TODO(b/160325686): Remove the batch logic after db migration.
     for (List<Key<PollMessage>> keys : Iterables.partition(query, BATCH_SIZE)) {
-      tm()
+      tm().transact(
-          .transact(
               () -> {
                 // Load poll messages and filter to just those of interest.
                 ImmutableList<PollMessage> pollMessages =
                     ofy().load().keys(keys).values().stream()
-                        .filter(pm -> pm instanceof OneTime)
                         .filter(pm -> isNullOrEmpty(message) || pm.getMsg().contains(message))
                         .collect(toImmutableList());
                 if (!dryRun) {
-                  ofy().delete().entities(pollMessages).now();
+                  pollMessages.forEach(PollFlowUtils::ackPollMessage);
                 }
                 pollMessages.forEach(
                     pm ->

core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java

@@ -14,10 +14,12 @@
 
 package google.registry.tools;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.persistResource;
 
+import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.Ofy;
@@ -25,6 +27,7 @@ import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessage.Autorenew;
 import google.registry.model.poll.PollMessage.OneTime;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
 import org.joda.time.DateTime;
@@ -47,12 +50,18 @@ public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesC
 
   @Test
   public void testSuccess_doesntDeletePollMessagesInFuture() throws Exception {
-    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
+    VKey<OneTime> pm1 =
-    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
+        persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey();
-    OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger");
+    VKey<OneTime> pm2 =
-    OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
+        persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives").createVKey();
+    VKey<OneTime> pm3 =
+        persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger").createVKey();
+    OneTime futurePollMessage =
+        persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
+    VKey<OneTime> pm4 = futurePollMessage.createVKey();
     runCommand("-c", "TheRegistrar");
-    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values()).containsExactly(pm4);
+    assertThat(tm().load(ImmutableList.of(pm1, pm2, pm3, pm4)).values())
+        .containsExactly(futurePollMessage);
     assertInStdout(
         "1-FSDGS-TLD-2406-624-2013,2013-05-01T22:33:44.000Z,ninelives",
         "1-FSDGS-TLD-2406-316-2014,2014-01-01T22:33:44.000Z,foobar",
@@ -61,10 +70,12 @@ public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesC
   }
 
   @Test
-  public void testSuccess_doesntDeleteAutorenewPollMessages() throws Exception {
+  public void testSuccess_resavesAutorenewPollMessages() throws Exception {
-    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
+    VKey<OneTime> pm1 =
-    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
+        persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey();
-    Autorenew pm3 =
+    VKey<OneTime> pm2 =
+        persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives").createVKey();
+    Autorenew autorenew =
         persistResource(
             new PollMessage.Autorenew.Builder()
                 .setId(624L)
@@ -73,26 +84,73 @@ public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesC
                         Key.create(DomainBase.class, "AAFSGS-TLD"), HistoryEntry.class, 99406L))
                 .setEventTime(DateTime.parse("2011-04-15T22:33:44Z"))
                 .setClientId("TheRegistrar")
+                .setMsg("autorenew")
                 .build());
+    Autorenew resaved =
+        autorenew.asBuilder().setEventTime(DateTime.parse("2012-04-15T22:33:44Z")).build();
+    VKey<Autorenew> pm3 = autorenew.createVKey();
     runCommand("-c", "TheRegistrar");
-    assertThat(ofy().load().entities(pm1, pm2, pm3).values()).containsExactly(pm3);
+    assertThat(tm().load(ImmutableList.of(pm1, pm2, pm3)).values()).containsExactly(resaved);
+    assertInStdout(
+        "1-AAFSGS-TLD-99406-624-2011,2011-04-15T22:33:44.000Z,autorenew",
+        "1-FSDGS-TLD-2406-624-2013,2013-05-01T22:33:44.000Z,ninelives",
+        "1-FSDGS-TLD-2406-316-2014,2014-01-01T22:33:44.000Z,foobar");
+  }
+
+  @Test
+  public void testSuccess_deletesExpiredAutorenewPollMessages() throws Exception {
+    VKey<OneTime> pm1 =
+        persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey();
+    VKey<OneTime> pm2 =
+        persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives").createVKey();
+    Autorenew autorenew =
+        persistResource(
+            new PollMessage.Autorenew.Builder()
+                .setId(624L)
+                .setParentKey(
+                    Key.create(
+                        Key.create(DomainBase.class, "AAFSGS-TLD"), HistoryEntry.class, 99406L))
+                .setEventTime(DateTime.parse("2011-04-15T22:33:44Z"))
+                .setAutorenewEndTime(DateTime.parse("2012-01-01T22:33:44Z"))
+                .setClientId("TheRegistrar")
+                .setMsg("autorenew")
+                .build());
+    VKey<Autorenew> pm3 = autorenew.createVKey();
+    runCommand("-c", "TheRegistrar");
+    assertThat(tm().load(ImmutableList.of(pm1, pm2, pm3))).isEmpty();
+    assertInStdout(
+        "1-AAFSGS-TLD-99406-624-2011,2011-04-15T22:33:44.000Z,autorenew",
+        "1-FSDGS-TLD-2406-624-2013,2013-05-01T22:33:44.000Z,ninelives",
+        "1-FSDGS-TLD-2406-316-2014,2014-01-01T22:33:44.000Z,foobar");
   }
 
   @Test
   public void testSuccess_onlyDeletesPollMessagesMatchingMessage() throws Exception {
-    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good");
+    VKey<OneTime> pm1 =
-    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
+        persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good")
-    OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "mmmmmfood");
+            .createVKey();
-    OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "time flies");
+    OneTime notMatched1 =
+        persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
+    VKey<OneTime> pm2 = notMatched1.createVKey();
+    VKey<OneTime> pm3 =
+        persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "mmmmmfood").createVKey();
+    OneTime notMatched2 =
+        persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "time flies");
+    VKey<OneTime> pm4 = notMatched2.createVKey();
     runCommand("-c", "TheRegistrar", "-m", "food");
-    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values()).containsExactly(pm2, pm4);
+    assertThat(tm().load(ImmutableList.of(pm1, pm2, pm3, pm4)).values())
+        .containsExactly(notMatched1, notMatched2);
   }
 
   @Test
   public void testSuccess_onlyDeletesPollMessagesMatchingClientId() throws Exception {
-    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good");
+    VKey<OneTime> pm1 =
-    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
+        persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good")
-    OneTime pm3 =
+            .createVKey();
+    VKey<OneTime> pm2 =
+        persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad")
+            .createVKey();
+    OneTime notMatched =
         persistResource(
             new PollMessage.OneTime.Builder()
                 .setId(2474L)
@@ -103,8 +161,9 @@ public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesC
                 .setEventTime(DateTime.parse("2013-06-01T22:33:44Z"))
                 .setMsg("baaaahh")
                 .build());
+    VKey<OneTime> pm3 = notMatched.createVKey();
     runCommand("-c", "TheRegistrar");
-    assertThat(ofy().load().entities(pm1, pm2, pm3).values()).containsExactly(pm3);
+    assertThat(tm().load(ImmutableList.of(pm1, pm2, pm3)).values()).containsExactly(notMatched);
   }
 
   @Test
@@ -114,7 +173,12 @@ public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesC
     OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger");
     OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
     runCommand("-c", "TheRegistrar", "-d");
-    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values())
+    assertThat(
+            tm().load(
+                    ImmutableList.of(pm1, pm2, pm3, pm4).stream()
+                        .map(OneTime::createVKey)
+                        .collect(toImmutableList()))
+                .values())
         .containsExactly(pm1, pm2, pm3, pm4);
   }