Save indexes when replaying EppResources SQL->DS (#1300)

* Save indexes when replaying EppResources SQL->DS

We implement this similarly to how we implement the
beforeSqlSaveOnReplay callback in the other direction -- a
beforeDatastoreSaveOnReplay method that is called when replaying a
Mutation to Datastore. This means that the asynchronous replay will
create the relevant ForeignKeyIndex and EppResourceIndex objects for
EppResources saved when SQL is primary.

core/src/main/java/google/registry/model/EppResource.java

@@ -20,6 +20,7 @@ import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.union;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@@ -37,6 +38,8 @@ import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Index;
 import google.registry.config.RegistryConfig;
 import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.CommitLogManifest;
 import google.registry.model.transfer.TransferData;
 import google.registry.persistence.VKey;
@@ -219,6 +222,14 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   @Override
   public abstract Builder<?, ?> asBuilder();
 
+  /** Used when replaying from SQL to DS to populate the Datastore indexes. */
+  protected void saveIndexesToDatastore() {
+    ofyTm()
+        .putAll(
+            ForeignKeyIndex.create(this, getDeletionTime()),
+            EppResourceIndex.create(Key.create(this)));
+  }
+
   /** EppResources that are loaded via foreign keys should implement this marker interface. */
   public interface ForeignKeyedEppResource {}
 

core/src/main/java/google/registry/model/contact/ContactResource.java

@@ -66,6 +66,11 @@ public class ContactResource extends ContactBase
     return ContactBase.cloneContactProjectedAtTime(this, now);
   }
 
+  @Override
+  public void beforeDatastoreSaveOnReplay() {
+    saveIndexesToDatastore();
+  }
+
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));

core/src/main/java/google/registry/model/domain/DomainBase.java

@@ -163,6 +163,11 @@ public class DomainBase extends DomainContent
     return cloneDomainProjectedAtTime(this, now);
   }
 
+  @Override
+  public void beforeDatastoreSaveOnReplay() {
+    saveIndexesToDatastore();
+  }
+
   public static VKey<DomainBase> createVKey(Key<DomainBase> key) {
     return VKey.create(DomainBase.class, key.getName(), key);
   }

core/src/main/java/google/registry/model/host/HostResource.java

@@ -51,6 +51,11 @@ public class HostResource extends HostBase
     return VKey.create(HostResource.class, getRepoId(), Key.create(this));
   }
 
+  @Override
+  public void beforeDatastoreSaveOnReplay() {
+    saveIndexesToDatastore();
+  }
+
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));

core/src/main/java/google/registry/model/replay/DatastoreEntity.java

@@ -27,4 +27,7 @@ import java.util.Optional;
 public interface DatastoreEntity {
 
   Optional<SqlEntity> toSqlEntity();
+
+  /** A method called before the object is saved to Datastore in asynchronous replay. */
+  default void beforeDatastoreSaveOnReplay() {}
 }

core/src/main/java/google/registry/model/replay/SqlEntity.java

@@ -29,10 +29,10 @@ public interface SqlEntity {
 
   Optional<DatastoreEntity> toDatastoreEntity();
 
-  /** A method that will ber called before the object is saved to SQL in asynchronous replay. */
+  /** A method that will be called before the object is saved to SQL in asynchronous replay. */
   default void beforeSqlSaveOnReplay() {}
 
-  /* Returns this entity's primary key field(s) in a string. */
+  /** Returns this entity's primary key field(s) in a string. */
   default String getPrimaryKeyString() {
     return jpaTm()
         .transact(

core/src/main/java/google/registry/persistence/transaction/Transaction.java

@@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
+import google.registry.model.replay.DatastoreEntity;
 import google.registry.model.replay.SqlEntity;
 import google.registry.persistence.VKey;
 import java.io.ByteArrayInputStream;
@@ -237,6 +238,10 @@ public class Transaction extends ImmutableObject implements Buildable {
 
     @Override
     public void writeToDatastore() {
+      // this should always be the case, but check just in case
+      if (entity instanceof DatastoreEntity) {
+        ((DatastoreEntity) entity).beforeDatastoreSaveOnReplay();
+      }
       ofyTm().put(entity);
     }
 

core/src/test/java/google/registry/model/contact/ContactResourceTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.testing.ContactResourceSubject.assertAboutContacts;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -29,6 +30,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
 import google.registry.model.ImmutableObjectSubject;
 import google.registry.model.contact.Disclose.PostalInfoChoice;
@@ -37,6 +39,9 @@ import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.PresenceMarker;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
 import google.registry.model.transfer.ContactTransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
@@ -277,4 +282,28 @@ public class ContactResourceTest extends EntityTestCase {
     // If there are circular references, this will overflow the stack.
     contactResource.toHydratedString();
   }
+
+  @Test
+  void testBeforeDatastoreSaveOnReplay_indexes() {
+    ImmutableList<ForeignKeyContactIndex> foreignKeyIndexes =
+        ofyTm().loadAllOf(ForeignKeyContactIndex.class);
+    ImmutableList<EppResourceIndex> eppResourceIndexes = ofyTm().loadAllOf(EppResourceIndex.class);
+    fakeClock.advanceOneMilli();
+    ofyTm()
+        .transact(
+            () -> {
+              foreignKeyIndexes.forEach(ofyTm()::delete);
+              eppResourceIndexes.forEach(ofyTm()::delete);
+            });
+    assertThat(ofyTm().loadAllOf(ForeignKeyContactIndex.class)).isEmpty();
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class)).isEmpty();
+
+    ofyTm().transact(() -> contactResource.beforeDatastoreSaveOnReplay());
+
+    assertThat(ofyTm().loadAllOf(ForeignKeyContactIndex.class))
+        .containsExactly(
+            ForeignKeyIndex.create(contactResource, contactResource.getDeletionTime()));
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class))
+        .containsExactly(EppResourceIndex.create(Key.create(contactResource)));
+  }
 }

core/src/test/java/google/registry/model/domain/DomainBaseTest.java

@@ -19,6 +19,7 @@ import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
@@ -51,6 +52,9 @@ import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.host.HostResource;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.tld.Registry;
@@ -893,6 +897,29 @@ public class DomainBaseTest extends EntityTestCase {
                 new BillEventInfo(recurringBillKey, null)));
   }
 
+  @Test
+  void testBeforeDatastoreSaveOnReplay_indexes() {
+    ImmutableList<ForeignKeyDomainIndex> foreignKeyIndexes =
+        ofyTm().loadAllOf(ForeignKeyDomainIndex.class);
+    ImmutableList<EppResourceIndex> eppResourceIndexes = ofyTm().loadAllOf(EppResourceIndex.class);
+    fakeClock.advanceOneMilli();
+    ofyTm()
+        .transact(
+            () -> {
+              foreignKeyIndexes.forEach(ofyTm()::delete);
+              eppResourceIndexes.forEach(ofyTm()::delete);
+            });
+    assertThat(ofyTm().loadAllOf(ForeignKeyDomainIndex.class)).isEmpty();
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class)).isEmpty();
+
+    ofyTm().transact(() -> domain.beforeDatastoreSaveOnReplay());
+
+    assertThat(ofyTm().loadAllOf(ForeignKeyDomainIndex.class))
+        .containsExactly(ForeignKeyIndex.create(domain, domain.getDeletionTime()));
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class))
+        .containsExactly(EppResourceIndex.create(Key.create(domain)));
+  }
+
   static class BillEventInfo extends ImmutableObject {
     VKey<BillingEvent.Recurring> billingEventRecurring;
     Long billingEventRecurringHistoryId;

core/src/test/java/google/registry/model/host/HostResourceTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -30,11 +31,15 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
+import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
 import google.registry.model.ImmutableObjectSubject;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.testing.DualDatabaseTest;
@@ -281,4 +286,27 @@ class HostResourceTest extends EntityTestCase {
     host = host.asBuilder().setLastTransferTime(day1).setLastSuperordinateChange(day2).build();
     assertThat(host.computeLastTransferTime(domain)).isEqualTo(day3);
   }
+
+  @TestOfyOnly
+  void testBeforeDatastoreSaveOnReplay_indexes() {
+    ImmutableList<ForeignKeyHostIndex> foreignKeyIndexes =
+        ofyTm().loadAllOf(ForeignKeyHostIndex.class);
+    ImmutableList<EppResourceIndex> eppResourceIndexes = ofyTm().loadAllOf(EppResourceIndex.class);
+    fakeClock.advanceOneMilli();
+    ofyTm()
+        .transact(
+            () -> {
+              foreignKeyIndexes.forEach(ofyTm()::delete);
+              eppResourceIndexes.forEach(ofyTm()::delete);
+            });
+    assertThat(ofyTm().loadAllOf(ForeignKeyHostIndex.class)).isEmpty();
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class)).isEmpty();
+
+    ofyTm().transact(() -> host.beforeDatastoreSaveOnReplay());
+
+    assertThat(ofyTm().loadAllOf(ForeignKeyHostIndex.class))
+        .containsExactly(ForeignKeyIndex.create(host, host.getDeletionTime()));
+    assertThat(ofyTm().loadAllOf(EppResourceIndex.class))
+        .containsExactly(EppResourceIndex.create(Key.create(host)));
+  }
 }

core/src/test/java/google/registry/model/replay/ReplicateToDatastoreActionTest.java

@@ -23,20 +23,16 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.testing.TestLogHandler;
-import com.googlecode.objectify.Key;
-import com.googlecode.objectify.annotation.Entity;
-import com.googlecode.objectify.annotation.Id;
-import google.registry.model.ImmutableObject;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.Ofy;
-import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestObject;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -56,8 +52,8 @@ public class ReplicateToDatastoreActionTest {
   public final AppEngineExtension appEngine =
       AppEngineExtension.builder()
           .withDatastoreAndCloudSql()
-          .withOfyTestEntities(TestEntity.class)
-          .withJpaUnitTestEntities(TestEntity.class)
+          .withOfyTestEntities(TestObject.class)
+          .withJpaUnitTestEntities(TestObject.class)
           .withClock(fakeClock)
           .build();
 
@@ -67,7 +63,7 @@ public class ReplicateToDatastoreActionTest {
   private final TestLogHandler logHandler = new TestLogHandler();
 
   @BeforeEach
-  public void setUp() {
+  void setUp() {
     injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
     // Use a single bucket to expose timestamp inversion problems.
     injectExtension.setStaticField(
@@ -76,19 +72,20 @@ public class ReplicateToDatastoreActionTest {
     DatabaseHelper.setMigrationScheduleToSqlPrimary(fakeClock);
     Logger.getLogger(ReplicateToDatastoreAction.class.getCanonicalName()).addHandler(logHandler);
     fakeClock.advanceBy(Duration.standardDays(1));
+    TestObject.beforeDatastoreSaveCallCount = 0;
   }
 
   @AfterEach
-  public void tearDown() {
+  void tearDown() {
     DatabaseHelper.removeDatabaseMigrationSchedule();
     fakeClock.disableAutoIncrement();
   }
 
   @RetryingTest(4)
-  public void testReplication() {
-    TestEntity foo = new TestEntity("foo");
-    TestEntity bar = new TestEntity("bar");
-    TestEntity baz = new TestEntity("baz");
+  void testReplication() {
+    TestObject foo = TestObject.create("foo");
+    TestObject bar = TestObject.create("bar");
+    TestObject baz = TestObject.create("baz");
 
     jpaTm()
         .transact(
@@ -115,9 +112,9 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
-  public void testReplayFromLastTxn() {
-    TestEntity foo = new TestEntity("foo");
-    TestEntity bar = new TestEntity("bar");
+  void testReplayFromLastTxn() {
+    TestObject foo = TestObject.create("foo");
+    TestObject bar = TestObject.create("bar");
 
     // Write a transaction containing "foo".
     jpaTm().transact(() -> jpaTm().insert(foo));
@@ -137,9 +134,9 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
-  public void testUnintentionalConcurrency() {
-    TestEntity foo = new TestEntity("foo");
-    TestEntity bar = new TestEntity("bar");
+  void testUnintentionalConcurrency() {
+    TestObject foo = TestObject.create("foo");
+    TestObject bar = TestObject.create("bar");
 
     // Write a transaction and run just the batch fetch.
     jpaTm().transact(() -> jpaTm().insert(foo));
@@ -172,9 +169,9 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
-  public void testMissingTransactions() {
+  void testMissingTransactions() {
     // Write a transaction (should have a transaction id of 1).
-    TestEntity foo = new TestEntity("foo");
+    TestObject foo = TestObject.create("foo");
     jpaTm().transact(() -> jpaTm().insert(foo));
 
     // Force the last transaction id back to -1 so that we look for transaction 0.
@@ -190,6 +187,15 @@ public class ReplicateToDatastoreActionTest {
             "Missing transaction: last transaction id = -1, next available transaction = 1");
   }
 
+  @Test
+  void testBeforeDatastoreSaveCallback() {
+    TestObject testObject = TestObject.create("foo");
+    jpaTm().transact(() -> jpaTm().put(testObject));
+    task.run();
+    assertThat(ofyTm().loadAllOf(TestObject.class)).containsExactly(testObject);
+    assertThat(TestObject.beforeDatastoreSaveCallCount).isEqualTo(1);
+  }
+
   @Test
   void testNotInMigrationState_doesNothing() {
     // set a schedule that backtracks the current status to DATASTORE_PRIMARY_READ_ONLY
@@ -208,10 +214,10 @@ public class ReplicateToDatastoreActionTest {
                         .build()));
     fakeClock.advanceBy(Duration.standardDays(1));
 
-    jpaTm().transact(() -> jpaTm().insert(new TestEntity("foo")));
+    jpaTm().transact(() -> jpaTm().insert(TestObject.create("foo")));
     task.run();
     // Replication shouldn't have happened
-    assertThat(ofyTm().loadAllOf(TestEntity.class)).isEmpty();
+    assertThat(ofyTm().loadAllOf(TestObject.class)).isEmpty();
     assertAboutLogs()
         .that(logHandler)
         .hasLogAtLevelWithMessage(
@@ -219,20 +225,4 @@ public class ReplicateToDatastoreActionTest {
             "Skipping ReplicateToDatastoreAction because we are in migration phase "
                 + "DATASTORE_PRIMARY_READ_ONLY.");
   }
-
-  @Entity(name = "ReplicationTestEntity")
-  @javax.persistence.Entity(name = "TestEntity")
-  private static class TestEntity extends ImmutableObject {
-    @Id @javax.persistence.Id private String name;
-
-    private TestEntity() {}
-
-    private TestEntity(String name) {
-      this.name = name;
-    }
-
-    public VKey<TestEntity> key() {
-      return VKey.create(TestEntity.class, name, Key.create(this));
-    }
-  }
 }

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

@@ -36,6 +36,7 @@ public class TestObject extends ImmutableObject implements DatastoreAndSqlEntity
 
   public static int beforeSqlSaveCallCount;
   public static int beforeSqlDeleteCallCount;
+  public static int beforeDatastoreSaveCallCount;
 
   @Parent @Transient Key<EntityGroupRoot> parent;
 
@@ -51,6 +52,10 @@ public class TestObject extends ImmutableObject implements DatastoreAndSqlEntity
     return field;
   }
 
+  public VKey<TestObject> key() {
+    return VKey.create(TestObject.class, id, Key.create(this));
+  }
+
   public static VKey<TestObject> createVKey(Key<TestObject> key) {
     return VKey.create(TestObject.class, key.getName(), key);
   }
@@ -80,6 +85,11 @@ public class TestObject extends ImmutableObject implements DatastoreAndSqlEntity
     beforeSqlSaveCallCount++;
   }
 
+  @Override
+  public void beforeDatastoreSaveOnReplay() {
+    beforeDatastoreSaveCallCount++;
+  }
+
   /** A test @VirtualEntity model object, which should not be persisted. */
   @Entity
   @VirtualEntity