Convert RefreshDnsOnHostRenameAction to tm (#1053)

* Convert RefreshDnsOnHostRenameAction to tm

This is not quite complete because it also requires the conversion of a
map-reduce which is in scope for an entirely different work.  Tests of the
map-reduce functionality are excluded from the SQL run.

This also requires the following additional fixes:

-  Convert Lock to tm, as doing so was necessary to get this action to work.
   As Lock is being targeted as DatastoreOnly, we convert all calls in it to
   use ofyTm()
-  Fix a bug in DualDatabaseTest (the check for an AppEngineExtension field is
   wrong, and captures fields of type Object as AppEngineExtension's)
-  Introduce another VKey.from() method that creates a VKey from a stringified
   Ofy Key.

* Rename VKey.from(String) to fromWebsafeKey

* Throw NoSuchElementE. instead of NPE

core/src/main/java/google/registry/batch/RefreshDnsOnHostRenameAction.java

@@ -25,7 +25,7 @@ import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.EppResourceUtils.isDeleted;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.latestOf;
 import static java.util.concurrent.TimeUnit.DAYS;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -44,7 +44,6 @@ import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
-import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskMetrics.OperationResult;
 import google.registry.dns.DnsQueue;
 import google.registry.mapreduce.MapreduceRunner;
@@ -64,6 +63,7 @@ import google.registry.util.SystemClock;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.logging.Level;
 import javax.annotation.Nullable;
@@ -123,7 +123,7 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
     }
 
     ImmutableList.Builder<DnsRefreshRequest> requestsBuilder = new ImmutableList.Builder<>();
-    ImmutableList.Builder<Key<HostResource>> hostKeys = new ImmutableList.Builder<>();
+    ImmutableList.Builder<VKey<HostResource>> hostKeys = new ImmutableList.Builder<>();
     final List<DnsRefreshRequest> requestsToDelete = new ArrayList<>();
 
     for (TaskHandle task : tasks) {
@@ -204,10 +204,10 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
         emit(true, true);
         return;
       }
-      Key<HostResource> referencingHostKey = null;
+      VKey<HostResource> referencingHostKey = null;
       for (DnsRefreshRequest request : refreshRequests) {
         if (isActive(domain, request.lastUpdateTime())
-            && domain.getNameservers().contains(VKey.from(request.hostKey()))) {
+            && domain.getNameservers().contains(request.hostKey())) {
           referencingHostKey = request.hostKey();
           break;
         }
@@ -293,7 +293,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     private static final long serialVersionUID = 1772812852271288622L;
 
-    abstract Key<HostResource> hostKey();
+    abstract VKey<HostResource> hostKey();
+
     abstract DateTime lastUpdateTime();
     abstract DateTime requestedTime();
     abstract boolean isRefreshNeeded();
@@ -301,7 +302,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     @AutoValue.Builder
     abstract static class Builder {
-      abstract Builder setHostKey(Key<HostResource> hostKey);
+      abstract Builder setHostKey(VKey<HostResource> hostKey);
+
       abstract Builder setLastUpdateTime(DateTime lastUpdateTime);
       abstract Builder setRequestedTime(DateTime requestedTime);
       abstract Builder setIsRefreshNeeded(boolean isRefreshNeeded);
@@ -314,10 +316,12 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
      */
     static DnsRefreshRequest createFromTask(TaskHandle task, DateTime now) throws Exception {
       ImmutableMap<String, String> params = ImmutableMap.copyOf(task.extractParams());
-      Key<HostResource> hostKey =
+      VKey<HostResource> hostKey =
-          Key.create(checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
+          VKey.fromWebsafeKey(
+              checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
       HostResource host =
-          checkNotNull(ofy().load().key(hostKey).now(), "Host to refresh doesn't exist");
+          tm().transact(() -> tm().loadByKeyIfPresent(hostKey))
+              .orElseThrow(() -> new NoSuchElementException("Host to refresh doesn't exist"));
       boolean isHostDeleted =
           isDeleted(host, latestOf(now, host.getUpdateTimestamp().getTimestamp()));
       if (isHostDeleted) {

core/src/main/java/google/registry/model/server/Lock.java

@@ -15,19 +15,20 @@
 package google.registry.model.server;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.google.common.flogger.FluentLogger;
+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.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
+import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.util.RequestStatusChecker;
 import google.registry.util.RequestStatusCheckerImpl;
@@ -190,12 +191,17 @@ public class Lock extends ImmutableObject implements DatastoreOnlyEntity, Serial
     // access to resources like GCS that can't be transactionally rolled back. Therefore, the lock
     // must be definitively acquired before it is used, even when called inside another transaction.
     AcquireResult acquireResult =
-        tm().transactNew(
+        ofyTm()
+            .transactNew(
                 () -> {
-                  DateTime now = tm().getTransactionTime();
+                  DateTime now = ofyTm().getTransactionTime();
 
                   // Checking if an unexpired lock still exists - if so, the lock can't be acquired.
-                  Lock lock = ofy().load().type(Lock.class).id(lockId).now();
+                  Lock lock =
+                      ofyTm()
+                          .loadByKeyIfPresent(
+                              VKey.createOfy(Lock.class, Key.create(Lock.class, lockId)))
+                          .orElse(null);
                   if (lock != null) {
                     logger.atInfo().log(
                         "Loaded existing lock: %s for request: %s", lock.lockId, lock.requestLogId);
@@ -218,7 +224,7 @@ public class Lock extends ImmutableObject implements DatastoreOnlyEntity, Serial
                   // Locks are not parented under an EntityGroupRoot (so as to avoid write
                   // contention) and
                   // don't need to be backed up.
-                  ofy().saveWithoutBackup().entity(newLock);
+                  ofyTm().putWithoutBackup(newLock);
 
                   return AcquireResult.create(now, lock, newLock, lockState);
                 });
@@ -231,21 +237,26 @@ public class Lock extends ImmutableObject implements DatastoreOnlyEntity, Serial
   /** Release the lock. */
   public void release() {
     // Just use the default clock because we aren't actually doing anything that will use the clock.
-    tm().transact(
+    ofyTm()
+        .transact(
             () -> {
               // To release a lock, check that no one else has already obtained it and if not
               // delete it. If the lock in Datastore was different then this lock is gone already;
               // this can happen if release() is called around the expiration time and the lock
               // expires underneath us.
-              Lock loadedLock = ofy().load().type(Lock.class).id(lockId).now();
+              Lock loadedLock =
+                  ofyTm()
+                      .loadByKeyIfPresent(
+                          VKey.createOfy(Lock.class, Key.create(Lock.class, lockId)))
+                      .orElse(null);
               if (Lock.this.equals(loadedLock)) {
                 // Use noBackupOfy() so that we don't create a commit log entry for deleting the
                 // lock.
                 logger.atInfo().log("Deleting lock: %s", lockId);
-                ofy().deleteWithoutBackup().entity(Lock.this);
+                ofyTm().deleteWithoutBackup(Lock.this);
 
                 lockMetrics.recordRelease(
-                    resourceName, tld, new Duration(acquiredTime, tm().getTransactionTime()));
+                    resourceName, tld, new Duration(acquiredTime, ofyTm().getTransactionTime()));
               } else {
                 logger.atSevere().log(
                     "The lock we acquired was transferred to someone else before we"

core/src/main/java/google/registry/persistence/VKey.java

@@ -223,4 +223,14 @@ public class VKey<T> extends ImmutableObject implements Serializable {
   public static <T> VKey<T> from(Key<T> key) {
     return VKeyTranslatorFactory.createVKey(key);
   }
+
+  /**
+   * Construct a VKey from the string representation of an ofy key.
+   *
+   * <p>TODO(b/184350590): After migration, we'll want remove the ofy key dependency from this.
+   */
+  @Nullable
+  public static <T> VKey<T> fromWebsafeKey(String ofyKeyRepr) {
+    return from(Key.create(ofyKeyRepr));
+  }
 }

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

@@ -19,7 +19,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.batch.AsyncTaskEnqueuer.QUEUE_ASYNC_HOST_RENAME;
 import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.newHostResource;
@@ -47,11 +47,14 @@ import google.registry.batch.AsyncTaskMetrics.OperationResult;
 import google.registry.batch.RefreshDnsOnHostRenameAction.RefreshDnsOnHostRenameReducer;
 import google.registry.model.host.HostResource;
 import google.registry.model.server.Lock;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.FakeSleeper;
 import google.registry.testing.InjectExtension;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import google.registry.util.AppEngineServiceUtils;
 import google.registry.util.RequestStatusChecker;
@@ -62,11 +65,11 @@ import java.util.Optional;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mock;
 
 /** Unit tests for {@link RefreshDnsOnHostRenameAction}. */
+@DualDatabaseTest
 public class RefreshDnsOnHostRenameActionTest
     extends MapreduceTestCase<RefreshDnsOnHostRenameAction> {
 
@@ -109,7 +112,7 @@ public class RefreshDnsOnHostRenameActionTest
     executeTasksUntilEmpty("mapreduce", clock);
     sleeper.sleep(millis(50));
     clock.advanceBy(standardSeconds(5));
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
   }
 
   /** Kicks off, but does not run, the mapreduce tasks. Useful for testing validation/setup. */
@@ -117,10 +120,13 @@ public class RefreshDnsOnHostRenameActionTest
     clock.advanceOneMilli();
     action.run();
     clock.advanceBy(standardSeconds(5));
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
   }
 
-  @Test
+  // TODO(b/181662306) None of the map reduce tests work with SQL since our map-reduce setup is
+  // inherently Datastore oriented, but this is a bigger task.
+
+  @TestOfyOnly
   void testSuccess_dnsUpdateEnqueued() throws Exception {
     HostResource host = persistActiveHost("ns1.example.tld");
     persistResource(newDomainBase("example.tld", host));
@@ -137,7 +143,7 @@ public class RefreshDnsOnHostRenameActionTest
     verifyNoMoreInteractions(action.asyncTaskMetrics);
   }
 
-  @Test
+  @TestOfyOnly
   void testSuccess_multipleHostsProcessedInBatch() throws Exception {
     HostResource host1 = persistActiveHost("ns1.example.tld");
     HostResource host2 = persistActiveHost("ns2.example.tld");
@@ -161,7 +167,7 @@ public class RefreshDnsOnHostRenameActionTest
     verifyNoMoreInteractions(action.asyncTaskMetrics);
   }
 
-  @Test
+  @TestOfyOnly
   void testSuccess_deletedHost_doesntTriggerDnsRefresh() throws Exception {
     HostResource host = persistDeletedHost("ns11.fakesss.tld", clock.nowUtc().minusDays(4));
     persistResource(newDomainBase("example1.tld", host));
@@ -176,7 +182,7 @@ public class RefreshDnsOnHostRenameActionTest
     verifyNoMoreInteractions(action.asyncTaskMetrics);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_noDnsTasksForDeletedDomain() throws Exception {
     HostResource renamedHost = persistActiveHost("ns1.example.tld");
     persistResource(
@@ -190,7 +196,7 @@ public class RefreshDnsOnHostRenameActionTest
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_hostDoesntExist_delaysTask() {
     HostResource host = newHostResource("ns1.example.tld");
     enqueuer.enqueueAsyncDnsRefresh(host, clock.nowUtc());
@@ -204,7 +210,7 @@ public class RefreshDnsOnHostRenameActionTest
     assertThat(acquireLock()).isPresent();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_cannotAcquireLock() {
     // Make lock acquisition fail.
     acquireLock();
@@ -213,7 +219,7 @@ public class RefreshDnsOnHostRenameActionTest
     assertNoDnsTasksEnqueued();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_mapreduceHasWorkToDo_lockIsAcquired() {
     HostResource host = persistActiveHost("ns1.example.tld");
     enqueuer.enqueueAsyncDnsRefresh(host, clock.nowUtc());
@@ -221,7 +227,7 @@ public class RefreshDnsOnHostRenameActionTest
     assertThat(acquireLock()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_noTasksToLease_releasesLockImmediately() {
     enqueueMapreduceOnly();
     assertNoDnsTasksEnqueued();

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

@@ -123,7 +123,7 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
       }
       fieldBuilder.addAll(
           Stream.of(clazz.getDeclaredFields())
-              .filter(field -> field.getType().isAssignableFrom(AppEngineExtension.class))
+              .filter(field -> AppEngineExtension.class.isAssignableFrom(field.getType()))
               .collect(toImmutableList()));
       return fieldBuilder.build();
     }