Make Caffeine cache loading work in non-GAE thread (#1634)

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

@@ -116,4 +116,9 @@ public class AppEngineEnvironment {
               }
             });
   }
+
+  /** Returns true if the current thread is in an App Engine Environment. */
+  public static boolean isInAppEngineEnvironment() {
+    return ApiProxy.getCurrentEnvironment() != null;
+  }
 }

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

@@ -18,9 +18,13 @@ import static com.google.common.base.Suppliers.memoizeWithExpiration;
 import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
+import com.github.benmanes.caffeine.cache.CacheLoader;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.common.base.Supplier;
+import google.registry.model.annotations.DeleteAfterMigration;
 import java.time.Duration;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Utility methods related to caching Datastore entities. */
 public class CacheUtils {
@@ -28,8 +32,8 @@ public class CacheUtils {
   /**
    * Memoize a supplier, with a short expiration specified in the environment config.
    *
-   * <p>Use this for things that might change while code is running. (For example, the various
+   * <p>Use this for things that might change while code is running. (For example, the various lists
-   * lists downloaded from the TMCH get updated in Datastore and the caches need to be refreshed.)
+   * downloaded from the TMCH get updated in Datastore and the caches need to be refreshed.)
    */
   public static <T> Supplier<T> memoizeWithShortExpiration(Supplier<T> original) {
     return tryMemoizeWithExpiration(getSingletonCacheRefreshDuration(), original);
@@ -73,4 +77,29 @@ public class CacheUtils {
     }
     return caffeine;
   }
+
+  /**
+   * A {@link CacheLoader} that automatically masquerade the background thread where the refresh
+   * action runs in to be an GAE thread.
+   */
+  @DeleteAfterMigration
+  public abstract static class AppEngineEnvironmentCacheLoader<K, V> implements CacheLoader<K, V> {
+
+    private static final AppEngineEnvironment environment = new AppEngineEnvironment();
+
+    @Override
+    public @Nullable V reload(@NonNull K key, @NonNull V oldValue) throws Exception {
+      V value;
+      boolean isMasqueraded = false;
+      if (!AppEngineEnvironment.isInAppEngineEnvironment()) {
+        environment.setEnvironmentForCurrentThread();
+        isMasqueraded = true;
+      }
+      value = load(key);
+      if (isMasqueraded) {
+        environment.unsetEnvironmentForCurrentThread();
+      }
+      return value;
+    }
+  }
 }

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

@@ -37,6 +37,7 @@ import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Index;
 import google.registry.config.RegistryConfig;
+import google.registry.model.CacheUtils.AppEngineEnvironmentCacheLoader;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
@@ -385,7 +386,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   }
 
   static final CacheLoader<VKey<? extends EppResource>, EppResource> CACHE_LOADER =
-      new CacheLoader<VKey<? extends EppResource>, EppResource>() {
+      new AppEngineEnvironmentCacheLoader<VKey<? extends EppResource>, EppResource>() {
 
         @Override
         public EppResource load(VKey<? extends EppResource> key) {

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -43,6 +43,7 @@ import com.googlecode.objectify.annotation.Index;
 import google.registry.config.RegistryConfig;
 import google.registry.model.BackupGroupRoot;
 import google.registry.model.CacheUtils;
+import google.registry.model.CacheUtils.AppEngineEnvironmentCacheLoader;
 import google.registry.model.EppResource;
 import google.registry.model.annotations.DeleteAfterMigration;
 import google.registry.model.annotations.ReportedOn;
@@ -256,7 +257,8 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   }
 
   static final CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
-      new CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
+      new AppEngineEnvironmentCacheLoader<
+          VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
 
         @Override
         public Optional<ForeignKeyIndex<?>> load(VKey<ForeignKeyIndex<?>> key) {

core/src/main/java/google/registry/model/tld/Registry.java

@@ -29,7 +29,6 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static org.joda.money.CurrencyUnit.USD;
 
-import com.github.benmanes.caffeine.cache.CacheLoader;
 import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
@@ -50,6 +49,7 @@ import com.googlecode.objectify.annotation.OnSave;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.CacheUtils;
+import google.registry.model.CacheUtils.AppEngineEnvironmentCacheLoader;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UnsafeSerializable;
@@ -268,7 +268,7 @@ public class Registry extends ImmutableObject
   private static final LoadingCache<String, Optional<Registry>> CACHE =
       CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration())
           .build(
-              new CacheLoader<String, Optional<Registry>>() {
+              new AppEngineEnvironmentCacheLoader<String, Optional<Registry>>() {
                 @Override
                 public Optional<Registry> load(final String tld) {
                   // Enter a transaction-less context briefly; we don't want to enroll every TLD in

core/src/main/java/google/registry/tmch/TmchCertificateAuthority.java

@@ -19,12 +19,12 @@ import static google.registry.config.RegistryConfig.ConfigModule.TmchCaMode.PROD
 import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 
-import com.github.benmanes.caffeine.cache.CacheLoader;
 import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.collect.ImmutableMap;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryConfig.ConfigModule.TmchCaMode;
 import google.registry.model.CacheUtils;
+import google.registry.model.CacheUtils.AppEngineEnvironmentCacheLoader;
 import google.registry.model.tmch.TmchCrl;
 import google.registry.util.Clock;
 import google.registry.util.X509Utils;
@@ -78,7 +78,7 @@ public final class TmchCertificateAuthority {
   private static final LoadingCache<TmchCaMode, X509CRL> CRL_CACHE =
       CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration())
           .build(
-              new CacheLoader<TmchCaMode, X509CRL>() {
+              new AppEngineEnvironmentCacheLoader<TmchCaMode, X509CRL>() {
                 @Override
                 public X509CRL load(final TmchCaMode tmchCaMode) throws GeneralSecurityException {
                   Optional<TmchCrl> storedCrl = TmchCrl.get();