diff --git a/core/src/main/java/google/registry/flows/poll/PollAckFlow.java b/core/src/main/java/google/registry/flows/poll/PollAckFlow.java index 77429e6ad..d027f166f 100644 --- a/core/src/main/java/google/registry/flows/poll/PollAckFlow.java +++ b/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; - 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; + boolean includeAckedMessageInCount = ackPollMessage(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. diff --git a/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java b/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java index 236711d3e..1e6db7ba5 100644 --- a/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java +++ b/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. + * + *

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; + } } diff --git a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java b/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java index b34cafb3c..b66208cc9 100644 --- a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java +++ b/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. * - *

This command only ACKs {@link OneTime} poll messages because that's our only use case so far, - * but it could be extended to ACK {@link Autorenew} poll messages as well if needed. The main + *

This command ACKs both {@link OneTime} and {@link Autorenew} poll messages. 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 query = getPollMessagesQuery(clientId, clock.nowUtc()).keys(); + // TODO(b/160325686): Remove the batch logic after db migration. for (List> keys : Iterables.partition(query, BATCH_SIZE)) { - tm() - .transact( + tm().transact( () -> { // Load poll messages and filter to just those of interest. ImmutableList 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 -> diff --git a/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java b/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java index 1d900bbb8..fe858465b 100644 --- a/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java +++ b/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 pm1 = + persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); + VKey pm2 = + persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives").createVKey(); + VKey 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 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 pm1 = + persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); + VKey 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 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 pm1 = + persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); + VKey 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 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"); - OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad"); - OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "mmmmmfood"); - OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "time flies"); + VKey pm1 = + persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good") + .createVKey(); + OneTime notMatched1 = + persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad"); + VKey pm2 = notMatched1.createVKey(); + VKey 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 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"); - OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad"); - OneTime pm3 = + VKey pm1 = + persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good") + .createVKey(); + VKey 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 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