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

This commit is contained in:
Lai Jiang 2022-05-18 22:01:44 -04:00 committed by GitHub
parent 126df845ea
commit 0144f332b0
6 changed files with 45 additions and 8 deletions

View file

@ -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;
}
}

View file

@ -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
* lists downloaded from the TMCH get updated in Datastore and the caches need to be refreshed.)
* <p>Use this for things that might change while code is running. (For example, the various lists
* 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;
}
}
}

View file

@ -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) {

View file

@ -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) {

View file

@ -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

View file

@ -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();