Fix create/update timestamp replay problems (#1515)

* Fix create/update timestamp replay problems

When CreateAutoTimestamp and UpdateAutoTimestamp are inserted into a
Transaction, their values are not populated in the same way as when they are
stored in the course of an SQL commit.  This results in different timestamp
values between SQL and datastore during the SQL -> DS replay.

Fix this by providing these values from the JPA transaction time when we're
doing transaction serialization.

This change also removes the initialization of the Ofy clock in
ExpandRecurringBillingEventsActionTest.  It's not necessary as the
ReplayExtension already takes care of this and doing it after the
ReplayExtension as we were breaks a test now that the update timestamps are
correct.

core/src/main/java/google/registry/model/translators/CreateAutoTimestampTranslatorFactory.java

@@ -14,12 +14,10 @@
 
 package google.registry.model.translators;
 
-import static com.google.common.base.MoreObjects.firstNonNull;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static org.joda.time.DateTimeZone.UTC;
 
 import google.registry.model.CreateAutoTimestamp;
-import google.registry.persistence.transaction.Transaction;
 import java.util.Date;
 import org.joda.time.DateTime;
 
@@ -47,13 +45,13 @@ public class CreateAutoTimestampTranslatorFactory
       /** Save a timestamp, setting it to the current time if it did not have a previous value. */
       @Override
       public Date saveValue(CreateAutoTimestamp pojoValue) {
-
-        // Don't do this if we're in the course of transaction serialization.
-        if (Transaction.inSerializationMode()) {
-          return pojoValue.getTimestamp() == null ? null : pojoValue.getTimestamp().toDate();
-        }
-
-        return firstNonNull(pojoValue.getTimestamp(), ofyTm().getTransactionTime()).toDate();
+        // Note that we use the current transaction manager -- we need to do this under JPA when we
+        // serialize the entity from a Transaction object, but we need to use the JPA transaction
+        // manager in that case.
+        return (pojoValue.getTimestamp() == null
+                ? tm().getTransactionTime()
+                : pojoValue.getTimestamp())
+            .toDate();
       }
     };
   }

core/src/main/java/google/registry/model/translators/UpdateAutoTimestampTranslatorFactory.java

@@ -14,6 +14,8 @@
 
 package google.registry.model.translators;
 
+import static com.google.common.base.Preconditions.checkState;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static org.joda.time.DateTimeZone.UTC;
 
@@ -48,9 +50,16 @@ public class UpdateAutoTimestampTranslatorFactory
       @Override
       public Date saveValue(UpdateAutoTimestamp pojoValue) {
 
-        // Don't do this if we're in the course of transaction serialization.
+        // If we're in the course of Transaction serialization, we have to use the transaction time
+        // here and the JPA transaction manager which is what will ultimately be saved during the
+        // commit.
+        // Note that this branch doesn't respect "auto update disabled", as this state is
+        // specifically to address replay, so we add a runtime check for this.
         if (Transaction.inSerializationMode()) {
-          return pojoValue.getTimestamp() == null ? null : pojoValue.getTimestamp().toDate();
+          checkState(
+              UpdateAutoTimestamp.autoUpdateEnabled(),
+              "Auto-update disabled during transaction serialization.");
+          return jpaTm().getTransactionTime().toDate();
         }
 
         return UpdateAutoTimestamp.autoUpdateEnabled()

core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -48,7 +48,6 @@ import google.registry.model.common.Cursor;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.Period;
-import google.registry.model.ofy.Ofy;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
@@ -56,7 +55,6 @@ import google.registry.model.tld.Registry;
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
-import google.registry.testing.InjectExtension;
 import google.registry.testing.ReplayExtension;
 import google.registry.testing.TestOfyAndSql;
 import google.registry.testing.TestOfyOnly;
@@ -78,11 +76,6 @@ public class ExpandRecurringBillingEventsActionTest
   private DateTime currentTestTime = DateTime.parse("1999-01-05T00:00:00Z");
   private final FakeClock clock = new FakeClock(currentTestTime);
 
-  @Order(Order.DEFAULT - 1)
-  @RegisterExtension
-  public final InjectExtension inject =
-      new InjectExtension().withStaticFieldOverride(Ofy.class, "clock", clock);
-
   @Order(Order.DEFAULT - 2)
   @RegisterExtension
   public final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -347,6 +347,12 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
             createBillingEvent));
     assertDnsTasksEnqueued(getUniqueIdFromCommand());
     assertEppResourceIndexEntityFor(domain);
+
+    replayExtension.expectUpdateFor(domain);
+
+    // Verify that all timestamps are correct after SQL -> DS replay.
+    // Added to confirm that timestamps get updated correctly.
+    replayExtension.enableDomainTimestampChecks();
   }
 
   private void assertNoLordn() throws Exception {
@@ -574,6 +580,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         ImmutableMap.of("DOMAIN", "otherexample.tld", "YEARS", "2"));
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "otherexample.tld")));
+    replayExtension.expectUpdateFor(reloadResourceByForeignKey());
   }
 
   @TestOfyAndSql
@@ -826,7 +833,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   @TestOfyAndSql
   void testSuccess_existedButWasDeleted() throws Exception {
     persistContactsAndHosts();
-    persistDeletedDomain(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
+    replayExtension.expectUpdateFor(
+        persistDeletedDomain(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1)));
     clock.advanceOneMilli();
     doSuccessfulTest();
   }

core/src/test/java/google/registry/flows/domain/DomainUpdateFlowTest.java

@@ -202,6 +202,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
             .setDomain(domain)
             .build());
     clock.advanceOneMilli();
+    replayExtension.expectUpdateFor(domain);
     return domain;
   }
 
@@ -212,9 +213,11 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
   private void doSuccessfulTest(String expectedXmlFilename) throws Exception {
     assertTransactionalFlow(true);
     runFlowAssertResponse(loadFile(expectedXmlFilename));
+    DomainBase domain = reloadResourceByForeignKey();
+    replayExtension.expectUpdateFor(domain);
     // Check that the domain was updated. These values came from the xml.
     assertAboutDomains()
-        .that(reloadResourceByForeignKey())
+        .that(domain)
         .hasStatusValue(StatusValue.CLIENT_HOLD)
         .and()
         .hasAuthInfoPwd("2BARfoo")
@@ -229,6 +232,10 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
     assertNoBillingEvents();
     assertDnsTasksEnqueued("example.tld");
     assertLastHistoryContainsResource(reloadResourceByForeignKey());
+
+    // Verify that all timestamps are correct after SQL -> DS replay.
+    // Added to confirm that timestamps get updated correctly.
+    replayExtension.enableDomainTimestampChecks();
   }
 
   @TestOfyAndSql
@@ -308,6 +315,8 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
     }
     persistResource(
         reloadResourceByForeignKey().asBuilder().setNameservers(nameservers.build()).build());
+    // Add a null update here so we don't compare.
+    replayExtension.expectUpdateFor(null);
     clock.advanceOneMilli();
   }
 

core/src/test/java/google/registry/testing/ReplayExtension.java

@@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.ReplayQueue;
 import google.registry.model.ofy.TransactionInfo;
@@ -41,6 +42,7 @@ import google.registry.persistence.transaction.Transaction.Update;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import javax.annotation.Nullable;
@@ -68,6 +70,8 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
   boolean inOfyContext;
   InjectExtension injectExtension = new InjectExtension();
   @Nullable ReplicateToDatastoreAction sqlToDsReplicator;
+  List<DomainBase> expectedUpdates = new ArrayList<>();
+  boolean enableDomainTimestampChecks;
 
   private ReplayExtension(
       FakeClock clock, boolean compare, @Nullable ReplicateToDatastoreAction sqlToDsReplicator) {
@@ -108,6 +112,30 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
     }
   }
 
+  /**
+   * Enable checking of domain timestamps during replay.
+   *
+   * <p>This was added to facilitate testing of a very specific bug wherein create/update
+   * auto-timestamps serialized to the SQL -> DS Transaction table had different values from those
+   * actually stored in SQL.
+   *
+   * <p>In order to use this, you also need to use expectUpdateFor() to store the states of a
+   * DomainBase object at a given point in time.
+   */
+  public void enableDomainTimestampChecks() {
+    enableDomainTimestampChecks = true;
+  }
+
+  /**
+   * If we're doing domain time checks, add the current state of a domain to check against.
+   *
+   * <p>A null argument is a placeholder to deal with b/217952766. Basically it allows us to ignore
+   * one particular state in the sequence (where the timestamp is not what we expect it to be).
+   */
+  public void expectUpdateFor(@Nullable DomainBase domain) {
+    expectedUpdates.add(domain);
+  }
+
   @Override
   public void beforeEach(ExtensionContext context) {
     // Use a single bucket to expose timestamp inversion problems. This typically happens when
@@ -253,6 +281,21 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
         assertAboutImmutableObjects()
             .that(fromDatastore)
             .isEqualAcrossDatabases(fromTransactionEntity);
+
+        // Check DomainBase timestamps if appropriate.
+        if (enableDomainTimestampChecks && fromTransactionEntity instanceof DomainBase) {
+          DomainBase expectedDomain = expectedUpdates.remove(0);
+
+          // Just skip it if the expectedDomain is null.
+          if (expectedDomain == null) {
+            continue;
+          }
+
+          DomainBase domainEntity = (DomainBase) fromTransactionEntity;
+          assertThat(domainEntity.getCreationTime()).isEqualTo(expectedDomain.getCreationTime());
+          assertThat(domainEntity.getUpdateTimestamp())
+              .isEqualTo(expectedDomain.getUpdateTimestamp());
+        }
       } else {
         Delete delete = (Delete) mutation;
         VKey<?> key = delete.getKey();