From f9184a37d5044b0930ea1f2f5888ef275a575986 Mon Sep 17 00:00:00 2001
From: Ben McIlwain <mcilwain@google.com>
Date: Tue, 10 Aug 2021 14:47:59 -0400
Subject: [PATCH] Load DatabaseMigrationStateSchedule in a more performant way
 (#1273)

This performs a direct load-by-key (the most efficient Datastore operation),
rather than attempting to load all entities by type using an ancestor query. The
existing implementation is possibly more error-prone as well, and might be
responsible for the "cross-group transaction need to be explicitly specified"
error we're seeing.
---
 .../DatabaseMigrationStateSchedule.java       | 25 +++++++++++++------
 .../java/google/registry/model/ofy/Ofy.java   |  2 +-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java b/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java
index 46d30debc..6ef46bd12 100644
--- a/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java
+++ b/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java
@@ -15,6 +15,8 @@
 package google.registry.model.common;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
@@ -24,6 +26,7 @@ import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSortedMap;
+import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Mapify;
@@ -32,6 +35,7 @@ import google.registry.model.common.TimedTransitionProperty.TimeMapper;
 import google.registry.model.common.TimedTransitionProperty.TimedTransition;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import java.time.Duration;
+import java.util.Optional;
 import org.joda.time.DateTime;
 
 /**
@@ -224,13 +228,20 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton
   /** Loads the currently-set migration schedule from Datastore, or the default if none exists. */
   @VisibleForTesting
   static TimedTransitionProperty<MigrationState, MigrationStateTransition> getUncached() {
-    return ofyTm()
-        .transactNew(
-            () ->
-                ofyTm()
-                    .loadSingleton(DatabaseMigrationStateSchedule.class)
-                    .map(s -> s.migrationTransitions)
-                    .orElse(DEFAULT_TRANSITION_MAP));
+    return Optional.ofNullable(
+            auditedOfy()
+                .doTransactionless(
+                    () ->
+                        auditedOfy()
+                            .load()
+                            .key(
+                                Key.create(
+                                    getCrossTldKey(),
+                                    DatabaseMigrationStateSchedule.class,
+                                    CrossTldSingleton.SINGLETON_ID))
+                            .now()))
+        .map(s -> s.migrationTransitions)
+        .orElse(DEFAULT_TRANSITION_MAP);
   }
 
   /**
diff --git a/core/src/main/java/google/registry/model/ofy/Ofy.java b/core/src/main/java/google/registry/model/ofy/Ofy.java
index d7257f5f2..abd38f40c 100644
--- a/core/src/main/java/google/registry/model/ofy/Ofy.java
+++ b/core/src/main/java/google/registry/model/ofy/Ofy.java
@@ -325,7 +325,7 @@ public class Ofy {
   }
 
   /** Execute some work in a transactionless context. */
-  <R> R doTransactionless(Supplier<R> work) {
+  public <R> R doTransactionless(Supplier<R> work) {
     try {
       com.googlecode.objectify.ObjectifyService.push(
           com.googlecode.objectify.ObjectifyService.ofy().transactionless());