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
2026-09-18 22:14:23 +00:00 · 2021-02-14 09:58:08 -05:00
parent f86936788e
commit bcc1924b24
6 changed files with 381 additions and 184 deletions
@@ -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() {}
 }
@@ -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) {
-                  return DatabaseTransitionSchedule.get(TransitionId.valueOf(transitionId));
+                public Optional<DatabaseTransitionSchedule> load(TransitionId 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. */
@@ -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(
-          () -> {
-            SignedMarkRevocationList datastoreList = loadFromDatastore();
-            suppressExceptionUnlessInTest(
-                () -> {
-                  loadAndCompareCloudSqlList(datastoreList);
-                },
-                "Error comparing signed mark revocation lists.");
-            return datastoreList;
-          });
+      memoizeWithShortExpiration(SignedMarkRevocationListDao::load);

  /** 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.trySave(this);
+    SignedMarkRevocationListDao.save(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() {
@@ -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);
-
-  /** Returns the most recent revision of the {@link SignedMarkRevocationList}, from cache. */
-  public static Optional<SignedMarkRevocationList> getLatestRevisionCached() {
-    return CACHE.get();
+  /**
+   * Loads the {@link SignedMarkRevocationList}.
+   *
+   * <p>Loads the list from the specified primary database, and attempts to load from the secondary
+   * database. If the load the secondary database fails, or the list from the secondary database
+   * 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
-   * error will be logged but no exception will be thrown.
+   * Save the given {@link SignedMarkRevocationList}
   *
-   * <p>This method is used during the dual-write phase of database migration as Datastore is still
-   * the authoritative database.
+   * <p>Saves the list to the specified primary database, and attempts to save to the secondary
+   * 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) {
-    suppressExceptionUnlessInTest(
-        () -> {
-          SignedMarkRevocationListDao.save(signedMarkRevocationList);
-          logger.atInfo().log(
-              "Inserted %,d signed mark revocations into Cloud SQL.",
-              signedMarkRevocationList.revokes.size());
-        },
-        "Error inserting signed mark revocations into Cloud SQL.");
+  static void save(SignedMarkRevocationList signedMarkRevocationList) {
+    PrimaryDatabase primaryDatabase =
+        tm().transactNew(
+                () ->
+                    DatabaseMigrationUtils.getPrimaryDatabase(
+                        TransitionId.SIGNED_MARK_REVOCATION_LIST));
+    if (primaryDatabase.equals(DATASTORE)) {
+      saveToDatastore(signedMarkRevocationList.revokes, signedMarkRevocationList.creationTime);
+      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()));
+            });
  }
 }