Expand AckPollMessagesCommand to ack PollMessage.Autorenew (#647)

* Expand AckPollMessagesCommand to ack PollMessage.Autorenew

* Rebase on master and address comment

* Resolve comments
This commit is contained in:
Shicong Huang 2020-07-01 15:06:35 -04:00 committed by GitHub
parent 74d0cdce5b
commit cf224a7e5b
4 changed files with 133 additions and 48 deletions

View file

@ -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.

View file

@ -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;
}
}

View file

@ -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,
* but it could be extended to ACK {@link Autorenew} poll messages as well if needed. The main
* <p>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<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()
.transact(
tm().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 ->

View file

@ -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");
OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger");
OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
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();
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 {
OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
Autorenew pm3 =
public void testSuccess_resavesAutorenewPollMessages() 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)
@ -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");
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<OneTime> 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<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");
OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
OneTime pm3 =
VKey<OneTime> pm1 =
persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good")
.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);
}