Refactor SignedMarkRevocationListDao for easy primary database cutover (#943)

* Refactor SignedMarkRevocationListDao for easy primary database cutover

* Fix javadoc comments

* Use PrimaryDatabase enum

* format fix

* fix up tests

* Fix punctuation

* Remove unnecessary else ifs

* Fix error messages

* spell out class name

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

@@ -16,6 +16,9 @@ package google.registry.model;
 
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryEnvironment;
+import google.registry.model.common.DatabaseTransitionSchedule;
+import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
+import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
 
 /** Utility methods related to migrating dual-read/dual-write entities. */
 public class DatabaseMigrationUtils {
@@ -34,5 +37,12 @@ public class DatabaseMigrationUtils {
     }
   }
 
+  /** Gets the value for the database currently considered primary. */
+  public static PrimaryDatabase getPrimaryDatabase(TransitionId transitionId) {
+    return DatabaseTransitionSchedule.getCached(transitionId)
+        .map(DatabaseTransitionSchedule::getPrimaryDatabase)
+        .orElse(PrimaryDatabase.DATASTORE);
+  }
+
   private DatabaseMigrationUtils() {}
 }

core/src/main/java/google/registry/model/common/DatabaseTransitionSchedule.java

@@ -96,15 +96,15 @@ public class DatabaseTransitionSchedule extends ImmutableObject implements Datas
       TimedTransitionProperty.forMapify(PrimaryDatabase.DATASTORE, PrimaryDatabaseTransition.class);
 
   /** A cache that loads the {@link DatabaseTransitionSchedule} for a given id. */
-  private static final LoadingCache<String, Optional<DatabaseTransitionSchedule>> CACHE =
+  private static final LoadingCache<TransitionId, Optional<DatabaseTransitionSchedule>> CACHE =
       CacheBuilder.newBuilder()
           .expireAfterWrite(
               java.time.Duration.ofMillis(getSingletonCacheRefreshDuration().getMillis()))
           .build(
-              new CacheLoader<String, Optional<DatabaseTransitionSchedule>>() {
+              new CacheLoader<TransitionId, Optional<DatabaseTransitionSchedule>>() {
                 @Override
-                public Optional<DatabaseTransitionSchedule> load(String transitionId) {
+                public Optional<DatabaseTransitionSchedule> load(TransitionId transitionId) {
-                  return DatabaseTransitionSchedule.get(TransitionId.valueOf(transitionId));
+                  return DatabaseTransitionSchedule.get(transitionId);
                 }
               });
 
@@ -136,7 +136,7 @@ public class DatabaseTransitionSchedule extends ImmutableObject implements Datas
    * <p>WARNING: The schedule returned by this method could be up to 10 minutes out of date.
    */
   public static Optional<DatabaseTransitionSchedule> getCached(TransitionId id) {
-    return CACHE.getUnchecked(id.name());
+    return CACHE.getUnchecked(id);
   }
 
   /** Returns the schedule for a given id. */

core/src/main/java/google/registry/model/smd/SignedMarkRevocationList.java

@@ -15,24 +15,13 @@
 package google.registry.model.smd;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static com.google.common.collect.Iterables.isEmpty;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
-import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.MapDifference;
-import com.google.common.collect.Maps;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EmbedMap;
 import com.googlecode.objectify.annotation.Entity;
@@ -45,9 +34,7 @@ import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.schema.replay.NonReplicatedEntity;
-import google.registry.util.CollectionUtils;
 import java.util.Map;
-import java.util.Optional;
 import javax.persistence.CollectionTable;
 import javax.persistence.Column;
 import javax.persistence.ElementCollection;
@@ -116,16 +103,7 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
    * single {@link SignedMarkRevocationList} object.
    */
   private static final Supplier<SignedMarkRevocationList> CACHE =
-      memoizeWithShortExpiration(
+      memoizeWithShortExpiration(SignedMarkRevocationListDao::load);
-          () -> {
-            SignedMarkRevocationList datastoreList = loadFromDatastore();
-            suppressExceptionUnlessInTest(
-                () -> {
-                  loadAndCompareCloudSqlList(datastoreList);
-                },
-                "Error comparing signed mark revocation lists.");
-            return datastoreList;
-          });
 
   /** Return a single logical instance that combines all Datastore shards. */
   public static SignedMarkRevocationList get() {
@@ -159,98 +137,10 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
 
   /** Save this list to Datastore in sharded form and to Cloud SQL. Returns {@code this}. */
   public SignedMarkRevocationList save() {
-    saveToDatastore();
+    SignedMarkRevocationListDao.save(this);
-    SignedMarkRevocationListDao.trySave(this);
     return this;
   }
 
-  /** Loads the shards from Datastore and combines them into one list. */
-  private static SignedMarkRevocationList loadFromDatastore() {
-    return tm().transactNewReadOnly(
-            () -> {
-              Iterable<SignedMarkRevocationList> shards =
-                  ofy().load().type(SignedMarkRevocationList.class).ancestor(getCrossTldKey());
-              DateTime creationTime =
-                  isEmpty(shards)
-                      ? START_OF_TIME
-                      : checkNotNull(Iterables.get(shards, 0).creationTime, "creationTime");
-              ImmutableMap.Builder<String, DateTime> revokes = new ImmutableMap.Builder<>();
-              for (SignedMarkRevocationList shard : shards) {
-                revokes.putAll(shard.revokes);
-                checkState(
-                    creationTime.equals(shard.creationTime),
-                    "Inconsistent creation times: %s vs. %s",
-                    creationTime,
-                    shard.creationTime);
-              }
-              return create(creationTime, revokes.build());
-            });
-  }
-
-  /** Save this list to Datastore in sharded form. */
-  private SignedMarkRevocationList saveToDatastore() {
-    tm().transact(
-            () -> {
-              ofy()
-                  .deleteWithoutBackup()
-                  .keys(
-                      ofy()
-                          .load()
-                          .type(SignedMarkRevocationList.class)
-                          .ancestor(getCrossTldKey())
-                          .keys());
-              ofy()
-                  .saveWithoutBackup()
-                  .entities(
-                      CollectionUtils.partitionMap(revokes, SHARD_SIZE).stream()
-                          .map(
-                              shardRevokes -> {
-                                SignedMarkRevocationList shard = create(creationTime, shardRevokes);
-                                shard.id = allocateId();
-                                shard.isShard =
-                                    true; // Avoid the exception in disallowUnshardedSaves().
-                                return shard;
-                              })
-                          .collect(toImmutableList()));
-            });
-    return this;
-  }
-
-  private static void loadAndCompareCloudSqlList(SignedMarkRevocationList datastoreList) {
-    // Lifted with some modifications from ClaimsListShard
-    Optional<SignedMarkRevocationList> maybeCloudSqlList =
-        SignedMarkRevocationListDao.getLatestRevision();
-    if (maybeCloudSqlList.isPresent()) {
-      SignedMarkRevocationList cloudSqlList = maybeCloudSqlList.get();
-      MapDifference<String, DateTime> diff =
-          Maps.difference(datastoreList.revokes, cloudSqlList.revokes);
-      if (!diff.areEqual()) {
-        if (diff.entriesDiffering().size() > 10) {
-          String message =
-              String.format(
-                  "Unequal SM revocation lists detected, Cloud SQL list with revision id %d has %d"
-                      + " different records than the current Datastore list.",
-                  cloudSqlList.revisionId, diff.entriesDiffering().size());
-          throw new RuntimeException(message);
-        } else {
-          StringBuilder diffMessage = new StringBuilder("Unequal SM revocation lists detected:\n");
-          diff.entriesDiffering()
-              .forEach(
-                  (label, valueDiff) ->
-                      diffMessage.append(
-                          String.format(
-                              "SMD %s has key %s in Datastore and key %s in Cloud SQL.\n",
-                              label, valueDiff.leftValue(), valueDiff.rightValue())));
-          throw new RuntimeException(diffMessage.toString());
-        }
-      }
-    } else {
-      if (datastoreList.size() != 0) {
-        throw new RuntimeException("Signed mark revocation list in Cloud SQL is empty.");
-      }
-    }
-  }
-
   /** As a safety mechanism, fail if someone tries to save this class directly. */
   @OnSave
   void disallowUnshardedSaves() {

core/src/main/java/google/registry/model/smd/SignedMarkRevocationListDao.java

@@ -14,28 +14,142 @@
 
 package google.registry.model.smd;
 
-import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.Iterables.isEmpty;
 import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
+import static google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase.DATASTORE;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.smd.SignedMarkRevocationList.SHARD_SIZE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
-import com.google.common.base.Supplier;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.MapDifference;
+import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
+import google.registry.model.DatabaseMigrationUtils;
+import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
+import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
+import google.registry.util.CollectionUtils;
+import java.util.Map;
 import java.util.Optional;
 import javax.persistence.EntityManager;
+import org.joda.time.DateTime;
 
 public class SignedMarkRevocationListDao {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
-  private static final Supplier<Optional<SignedMarkRevocationList>> CACHE =
+  /**
-      memoizeWithShortExpiration(SignedMarkRevocationListDao::getLatestRevision);
+   * Loads the {@link SignedMarkRevocationList}.
-
+   *
-  /** Returns the most recent revision of the {@link SignedMarkRevocationList}, from cache. */
+   * <p>Loads the list from the specified primary database, and attempts to load from the secondary
-  public static Optional<SignedMarkRevocationList> getLatestRevisionCached() {
+   * database. If the load the secondary database fails, or the list from the secondary database
-    return CACHE.get();
+   * does not match the list from the primary database, the error will be logged but no exception
+   * will be thrown.
+   */
+  static SignedMarkRevocationList load() {
+    PrimaryDatabase primaryDatabase =
+        tm().transactNew(
+                () ->
+                    DatabaseMigrationUtils.getPrimaryDatabase(
+                        TransitionId.SIGNED_MARK_REVOCATION_LIST));
+    Optional<SignedMarkRevocationList> primaryList =
+        primaryDatabase.equals(DATASTORE) ? loadFromDatastore() : loadFromCloudSql();
+    if (!primaryList.isPresent()) {
+      throw new IllegalStateException(
+          String.format(
+              "SignedMarkRevocationList not found in the primary database (%s).",
+              primaryDatabase.name()));
+    }
+    suppressExceptionUnlessInTest(
+        () -> loadAndCompare(primaryDatabase, primaryList.get()),
+        String.format(
+            "Error loading and comparing the SignedMarkRevocationList from the secondary database"
+                + " (%s).",
+            primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore"));
+    return primaryList.get();
   }
 
-  public static Optional<SignedMarkRevocationList> getLatestRevision() {
+  /**
+   * Loads the list from the secondary database and compares it to the list from the primary
+   * database.
+   */
+  private static void loadAndCompare(
+      PrimaryDatabase primaryDatabase, SignedMarkRevocationList primaryList) {
+    Optional<SignedMarkRevocationList> secondaryList =
+        primaryDatabase.equals(DATASTORE) ? loadFromCloudSql() : loadFromDatastore();
+    if (secondaryList.isPresent() && !isNullOrEmpty(secondaryList.get().revokes)) {
+      MapDifference<String, DateTime> diff =
+          Maps.difference(primaryList.revokes, secondaryList.get().revokes);
+      if (!diff.areEqual()) {
+        if (diff.entriesDiffering().size() > 10) {
+          String message =
+              String.format(
+                  "Unequal SignedMarkRevocationList detected, %s list with revision id"
+                      + " %d has %d different records than the current primary database list.",
+                  primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore",
+                  secondaryList.get().revisionId,
+                  diff.entriesDiffering().size());
+          throw new IllegalStateException(message);
+        } else {
+          StringBuilder diffMessage =
+              new StringBuilder("Unequal SignedMarkRevocationList detected:\n");
+          diff.entriesDiffering()
+              .forEach(
+                  (label, valueDiff) ->
+                      diffMessage.append(
+                          String.format(
+                              "SMD %s has key %s in %s and key %s in secondary database.\n",
+                              label,
+                              valueDiff.leftValue(),
+                              primaryDatabase.name(),
+                              valueDiff.rightValue())));
+          throw new IllegalStateException(diffMessage.toString());
+        }
+      }
+    } else {
+      if (primaryList.size() != 0) {
+        throw new IllegalStateException(
+            String.format(
+                "SignedMarkRevocationList in %s is empty while it is not empty in the primary"
+                    + " database.",
+                primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore"));
+      }
+    }
+  }
+
+  /** Loads the shards from Datastore and combines them into one list. */
+  private static Optional<SignedMarkRevocationList> loadFromDatastore() {
+    return tm().transactNewReadOnly(
+            () -> {
+              Iterable<SignedMarkRevocationList> shards =
+                  ofy().load().type(SignedMarkRevocationList.class).ancestor(getCrossTldKey());
+              DateTime creationTime =
+                  isEmpty(shards)
+                      ? START_OF_TIME
+                      : checkNotNull(Iterables.get(shards, 0).creationTime, "creationTime");
+              ImmutableMap.Builder<String, DateTime> revokes = new ImmutableMap.Builder<>();
+              for (SignedMarkRevocationList shard : shards) {
+                revokes.putAll(shard.revokes);
+                checkState(
+                    creationTime.equals(shard.creationTime),
+                    "Inconsistent creation times in Datastore shard: %s vs. %s",
+                    creationTime,
+                    shard.creationTime);
+              }
+              return Optional.of(SignedMarkRevocationList.create(creationTime, revokes.build()));
+            });
+  }
+
+  private static Optional<SignedMarkRevocationList> loadFromCloudSql() {
     return jpaTm()
         .transact(
             () -> {
@@ -54,24 +168,65 @@ public class SignedMarkRevocationListDao {
   }
 
   /**
-   * Try to save the given {@link SignedMarkRevocationList} into Cloud SQL. If the save fails, the
+   * Save the given {@link SignedMarkRevocationList}
-   * error will be logged but no exception will be thrown.
    *
-   * <p>This method is used during the dual-write phase of database migration as Datastore is still
+   * <p>Saves the list to the specified primary database, and attempts to save to the secondary
-   * the authoritative database.
+   * database. If the save to the secondary database fails, the error will be logged but no
+   * exception will be thrown.
    */
-  static void trySave(SignedMarkRevocationList signedMarkRevocationList) {
+  static void save(SignedMarkRevocationList signedMarkRevocationList) {
-    suppressExceptionUnlessInTest(
+    PrimaryDatabase primaryDatabase =
-        () -> {
+        tm().transactNew(
-          SignedMarkRevocationListDao.save(signedMarkRevocationList);
+                () ->
-          logger.atInfo().log(
+                    DatabaseMigrationUtils.getPrimaryDatabase(
-              "Inserted %,d signed mark revocations into Cloud SQL.",
+                        TransitionId.SIGNED_MARK_REVOCATION_LIST));
-              signedMarkRevocationList.revokes.size());
+    if (primaryDatabase.equals(DATASTORE)) {
-        },
+      saveToDatastore(signedMarkRevocationList.revokes, signedMarkRevocationList.creationTime);
-        "Error inserting signed mark revocations into Cloud SQL.");
+      suppressExceptionUnlessInTest(
+          () -> SignedMarkRevocationListDao.saveToCloudSql(signedMarkRevocationList),
+          "Error inserting signed mark revocations into secondary database (Cloud SQL).");
+    } else {
+      SignedMarkRevocationListDao.saveToCloudSql(signedMarkRevocationList);
+      suppressExceptionUnlessInTest(
+          () ->
+              saveToDatastore(
+                  signedMarkRevocationList.revokes, signedMarkRevocationList.creationTime),
+          "Error inserting signed mark revocations into secondary database (Datastore).");
+    }
   }
 
-  private static void save(SignedMarkRevocationList signedMarkRevocationList) {
+  private static void saveToCloudSql(SignedMarkRevocationList signedMarkRevocationList) {
     jpaTm().transact(() -> jpaTm().getEntityManager().persist(signedMarkRevocationList));
+    logger.atInfo().log(
+        "Inserted %,d signed mark revocations into Cloud SQL.",
+        signedMarkRevocationList.revokes.size());
+  }
+
+  private static void saveToDatastore(Map<String, DateTime> revokes, DateTime creationTime) {
+    tm().transact(
+            () -> {
+              ofy()
+                  .deleteWithoutBackup()
+                  .keys(
+                      ofy()
+                          .load()
+                          .type(SignedMarkRevocationList.class)
+                          .ancestor(getCrossTldKey())
+                          .keys());
+              ofy()
+                  .saveWithoutBackup()
+                  .entities(
+                      CollectionUtils.partitionMap(revokes, SHARD_SIZE).stream()
+                          .map(
+                              shardRevokes -> {
+                                SignedMarkRevocationList shard =
+                                    SignedMarkRevocationList.create(creationTime, shardRevokes);
+                                shard.id = allocateId();
+                                shard.isShard =
+                                    true; // Avoid the exception in disallowUnshardedSaves().
+                                return shard;
+                              })
+                          .collect(toImmutableList()));
+            });
   }
 }

core/src/test/java/google/registry/model/smd/SignedMarkRevocationListDaoTest.java

@@ -16,21 +16,35 @@ package google.registry.model.smd;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.config.RegistryEnvironment;
+import google.registry.model.EntityTestCase;
+import google.registry.model.common.DatabaseTransitionSchedule;
+import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
+import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition;
+import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
+import google.registry.model.common.TimedTransitionProperty;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.DualDatabaseTest;
-import google.registry.testing.FakeClock;
+import google.registry.testing.SystemPropertyExtension;
+import google.registry.testing.TestOfyAndSql;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 @DualDatabaseTest
-public class SignedMarkRevocationListDaoTest {
+public class SignedMarkRevocationListDaoTest extends EntityTestCase {
-
-  private final FakeClock fakeClock = new FakeClock();
 
   @RegisterExtension
   final JpaIntegrationWithCoverageExtension jpa =
@@ -40,50 +54,186 @@ public class SignedMarkRevocationListDaoTest {
   @Order(value = 1)
   final DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
 
-  @Test
+  @RegisterExtension
-  void testSave_success() {
+  @Order(value = Integer.MAX_VALUE)
+  final SystemPropertyExtension systemPropertyExtension = new SystemPropertyExtension();
+
+  @BeforeEach
+  void setup() {
+    fakeClock.setTo(DateTime.parse("1984-12-21T00:00:00.000Z"));
+    DatabaseTransitionSchedule schedule =
+        DatabaseTransitionSchedule.create(
+            TransitionId.SIGNED_MARK_REVOCATION_LIST,
+            TimedTransitionProperty.fromValueMap(
+                ImmutableSortedMap.of(
+                    START_OF_TIME,
+                    PrimaryDatabase.DATASTORE,
+                    fakeClock.nowUtc().plusDays(1),
+                    PrimaryDatabase.CLOUD_SQL),
+                PrimaryDatabaseTransition.class));
+
+    tm().transactNew(() -> ofy().saveWithoutBackup().entity(schedule).now());
+  }
+
+  @TestOfyAndSql
+  void testSave_datastorePrimary_success() {
     SignedMarkRevocationList list =
         SignedMarkRevocationList.create(
             fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
-    SignedMarkRevocationListDao.trySave(list);
+    SignedMarkRevocationListDao.save(list);
-    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.getLatestRevision().get();
+    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.load();
+    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list, "revisionId");
+  }
+
+  @TestOfyAndSql
+  void testSave_cloudSqlPrimary_success() {
+    fakeClock.advanceBy(Duration.standardDays(5));
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.load();
     assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list);
   }
 
-  @Test
+  @TestOfyAndSql
-  void testRetrieval_notPresent() {
+  void testSaveAndLoad_datastorePrimary_emptyList() {
-    assertThat(SignedMarkRevocationListDao.getLatestRevision().isPresent()).isFalse();
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(START_OF_TIME, ImmutableMap.of());
+    SignedMarkRevocationListDao.save(list);
+    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.load();
+    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list, "revisionId");
   }
 
-  @Test
+  @TestOfyAndSql
-  void testSaveAndRetrieval_emptyList() {
+  void testSaveAndLoad_cloudSqlPrimary_emptyList() {
+    fakeClock.advanceBy(Duration.standardDays(5));
     SignedMarkRevocationList list =
         SignedMarkRevocationList.create(fakeClock.nowUtc(), ImmutableMap.of());
-    SignedMarkRevocationListDao.trySave(list);
+    SignedMarkRevocationListDao.save(list);
-    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.getLatestRevision().get();
+    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.load();
-    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list);
+    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list, "revisionId");
   }
 
-  @Test
+  @TestOfyAndSql
-  void testSave_multipleVersions() {
+  void testSave_datastorePrimary_multipleVersions() {
     SignedMarkRevocationList list =
         SignedMarkRevocationList.create(
             fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
-    SignedMarkRevocationListDao.trySave(list);
+    SignedMarkRevocationListDao.save(list);
-    assertThat(
+    assertThat(SignedMarkRevocationListDao.load().isSmdRevoked("mark", fakeClock.nowUtc()))
-            SignedMarkRevocationListDao.getLatestRevision()
-                .get()
-                .isSmdRevoked("mark", fakeClock.nowUtc()))
         .isTrue();
 
     // Now remove the revocation
     SignedMarkRevocationList secondList =
         SignedMarkRevocationList.create(fakeClock.nowUtc(), ImmutableMap.of());
-    SignedMarkRevocationListDao.trySave(secondList);
+    SignedMarkRevocationListDao.save(secondList);
-    assertThat(
+    assertThat(SignedMarkRevocationListDao.load().isSmdRevoked("mark", fakeClock.nowUtc()))
-            SignedMarkRevocationListDao.getLatestRevision()
-                .get()
-                .isSmdRevoked("mark", fakeClock.nowUtc()))
         .isFalse();
   }
+
+  @TestOfyAndSql
+  void testSave_cloudSqlPrimary_multipleVersions() {
+    fakeClock.advanceBy(Duration.standardDays(5));
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    assertThat(SignedMarkRevocationListDao.load().isSmdRevoked("mark", fakeClock.nowUtc()))
+        .isTrue();
+
+    // Now remove the revocation
+    SignedMarkRevocationList secondList =
+        SignedMarkRevocationList.create(fakeClock.nowUtc(), ImmutableMap.of());
+    SignedMarkRevocationListDao.save(secondList);
+    assertThat(SignedMarkRevocationListDao.load().isSmdRevoked("mark", fakeClock.nowUtc()))
+        .isFalse();
+  }
+
+  @TestOfyAndSql
+  void testLoad_datastorePrimary_unequalLists() {
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    SignedMarkRevocationList list2 =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(3)));
+    jpaTm().transact(() -> jpaTm().put(list2));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, SignedMarkRevocationListDao::load);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "SMD mark has key 1984-12-20T23:00:00.000Z in DATASTORE and key"
+                + " 1984-12-20T21:00:00.000Z in secondary database.");
+  }
+
+  @TestOfyAndSql
+  void testLoad_cloudSqlPrimary_unequalLists() {
+    fakeClock.advanceBy(Duration.standardDays(5));
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    SignedMarkRevocationList list2 =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(3)));
+    jpaTm().transact(() -> jpaTm().put(list2));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, SignedMarkRevocationListDao::load);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "SMD mark has key 1984-12-25T21:00:00.000Z in CLOUD_SQL and key"
+                + " 1984-12-25T23:00:00.000Z in secondary database.");
+  }
+
+  @TestOfyAndSql
+  void testLoad_cloudSqlPrimary_unequalLists_succeedsInProduction() {
+    RegistryEnvironment.PRODUCTION.setup(systemPropertyExtension);
+    fakeClock.advanceBy(Duration.standardDays(5));
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    SignedMarkRevocationList list2 =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(3)));
+    jpaTm().transact(() -> jpaTm().put(list2));
+    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.load();
+    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list2, "revisionId");
+  }
+
+  @TestOfyAndSql
+  void testLoad_datastorePrimary_noListInCloudSql() {
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    SignedMarkRevocationListDao.save(list);
+    jpaTm().transact(() -> jpaTm().delete(list));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, SignedMarkRevocationListDao::load);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "SignedMarkRevocationList in Cloud SQL is empty while it is not empty in the"
+                + " primary database.");
+  }
+
+  @TestOfyAndSql
+  void testLoad_cloudSqlPrimary_noListInDatastore() {
+    fakeClock.advanceBy(Duration.standardDays(5));
+    SignedMarkRevocationList list =
+        SignedMarkRevocationList.create(
+            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
+    jpaTm().transact(() -> jpaTm().put(list));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, SignedMarkRevocationListDao::load);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "SignedMarkRevocationList in Datastore is empty while it is not empty in the"
+                + " primary database.");
+  }
 }

core/src/test/java/google/registry/model/smd/SignedMarkRevocationListTest.java

@@ -139,18 +139,6 @@ public class SignedMarkRevocationListTest {
         .isEqualTo(DateTime.parse("2000-01-01T00:00:00Z"));
   }
 
-  @Test
-  void test_getCreationTime_missingInCloudSQL() {
-    clock.setTo(DateTime.parse("2000-01-01T00:00:00Z"));
-    createSaveGetHelper(1);
-    jpaTm().transact(() -> jpaTm().delete(SignedMarkRevocationListDao.getLatestRevision().get()));
-    RuntimeException thrown =
-        assertThrows(RuntimeException.class, () -> SignedMarkRevocationList.get());
-    assertThat(thrown)
-        .hasMessageThat()
-        .isEqualTo("Signed mark revocation list in Cloud SQL is empty.");
-  }
-
   @Test
   void test_getCreationTime_unequalListsInDatabases() {
     clock.setTo(DateTime.parse("2000-01-01T00:00:00Z"));
@@ -159,11 +147,15 @@ public class SignedMarkRevocationListTest {
     for (int i = 0; i < 3; i++) {
       revokes.put(Integer.toString(i), clock.nowUtc());
     }
-    SignedMarkRevocationListDao.trySave(
+    jpaTm()
-        SignedMarkRevocationList.create(clock.nowUtc(), revokes.build()));
+        .transact(
+            () ->
+                jpaTm()
+                    .getEntityManager()
+                    .persist(SignedMarkRevocationList.create(clock.nowUtc(), revokes.build())));
     RuntimeException thrown =
         assertThrows(RuntimeException.class, () -> SignedMarkRevocationList.get());
-    assertThat(thrown).hasMessageThat().contains("Unequal SM revocation lists detected:");
+    assertThat(thrown).hasMessageThat().contains("Unequal SignedMarkRevocationList detected:");
   }
 
   @Test