Replace caffeine cache with reference to full map

SingularityService/pom.xml

@@ -510,11 +510,6 @@
       <scope>test</scope>
     </dependency>
 
-    <dependency>
-      <groupId>com.github.ben-manes.caffeine</groupId>
-      <artifactId>caffeine</artifactId>
-    </dependency>
-
   </dependencies>
 
   <build>

SingularityService/src/main/java/com/hubspot/singularity/SingularityServiceModule.java

@@ -1,13 +1,10 @@
 package com.hubspot.singularity;
 
-import com.github.benmanes.caffeine.cache.Cache;
-import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.common.base.Function;
 import com.google.inject.Binder;
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
-import com.google.inject.name.Named;
 import com.hubspot.dropwizard.guicier.DropwizardAwareModule;
 import com.hubspot.mesos.client.SingularityMesosClientModule;
 import com.hubspot.mesos.client.UserAndPassword;
@@ -29,14 +26,10 @@
 import com.hubspot.singularity.resources.SingularityOpenApiResource;
 import com.hubspot.singularity.resources.SingularityResourceModule;
 import com.hubspot.singularity.scheduler.SingularitySchedulerModule;
-import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.TimeUnit;
 
 public class SingularityServiceModule
   extends DropwizardAwareModule<SingularityConfiguration> {
-  public static final String REQUESTS_CAFFEINE_CACHE =
-    "singularity.service.resources.request";
   private final Function<SingularityConfiguration, Module> dbModuleProvider;
   private Optional<Class<? extends LoadBalancerClient>> lbClientClass = Optional.empty();
 
@@ -135,16 +128,4 @@ public IndexViewConfiguration provideIndexViewConfiguration(
         .contains(SingularityAuthenticatorClass.WEBHOOK)
     );
   }
-
-  @Provides
-  @Singleton
-  @Named(REQUESTS_CAFFEINE_CACHE)
-  public Cache<String, List<SingularityRequestParent>> getRequestsCaffeineCache() {
-    SingularityConfiguration configuration = getConfiguration();
-
-    return Caffeine
-      .newBuilder()
-      .expireAfterWrite(configuration.getCaffeineCacheTtl(), TimeUnit.SECONDS)
-      .build();
-  }
 }

SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java

@@ -444,10 +444,12 @@ public class SingularityConfiguration extends Configuration {
   private double statusQueueNearlyFull = 0.8;
 
   // Enable caffeine cache on heavily requested endpoint
-  private boolean useCaffeineCache = false;
+  private boolean useApiCacheInRequestManager = false;
+  private boolean useApiCacheInDeployManager = false;
 
-  // Caffeine cache ttl
-  private int caffeineCacheTtl = 1;
+  // Atomic Reference cache TTLs
+  private int deployCacheTtlInSeconds = 1;
+  private int requestCacheTtlInSeconds = 1;
 
   public long getAskDriverToKillTasksAgainAfterMillis() {
     return askDriverToKillTasksAgainAfterMillis;
@@ -2076,19 +2078,35 @@ public void setStatusQueueNearlyFull(double statusQueueNearlyFull) {
     this.statusQueueNearlyFull = statusQueueNearlyFull;
   }
 
-  public boolean useCaffeineCache() {
-    return useCaffeineCache;
+  public boolean useApiCacheInRequestManager() {
+    return useApiCacheInRequestManager;
   }
 
-  public void setUseCaffeineCache(boolean useCaffeineCache) {
-    this.useCaffeineCache = useCaffeineCache;
+  public void setUseApiCacheInRequestManager(boolean useApiCacheInRequestManager) {
+    this.useApiCacheInRequestManager = useApiCacheInRequestManager;
   }
 
-  public int getCaffeineCacheTtl() {
-    return caffeineCacheTtl;
+  public boolean useApiCacheInDeployManager() {
+    return useApiCacheInDeployManager;
   }
 
-  public void setCaffeineCacheTtl(int caffeineCacheTtl) {
-    this.caffeineCacheTtl = caffeineCacheTtl;
+  public void setUseApiCacheInDeployManager(boolean useApiCacheInDeployManager) {
+    this.useApiCacheInDeployManager = useApiCacheInDeployManager;
+  }
+
+  public int getDeployCacheTtl() {
+    return deployCacheTtlInSeconds;
+  }
+
+  public void setDeployCacheTtl(int deployCacheTtlInSeconds) {
+    this.deployCacheTtlInSeconds = deployCacheTtlInSeconds;
+  }
+
+  public int getRequestCacheTtl() {
+    return requestCacheTtlInSeconds;
+  }
+
+  public void setRequestCacheTtl(int requestCacheTtlInSeconds) {
+    this.requestCacheTtlInSeconds = requestCacheTtlInSeconds;
   }
 }

SingularityService/src/main/java/com/hubspot/singularity/data/ApiCache.java

@@ -0,0 +1,77 @@
+package com.hubspot.singularity.data;
+
+import com.google.inject.Inject;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ApiCache<K, V> {
+  private static final Logger LOG = LoggerFactory.getLogger(ApiCache.class);
+
+  public final boolean isEnabled;
+  private final AtomicReference<Map<K, V>> zkValues;
+  private final Supplier<Map<K, V>> supplyMap;
+
+  @Inject
+  public ApiCache(
+    boolean isEnabled,
+    int cacheTtl,
+    Supplier<Map<K, V>> supplyMap,
+    ScheduledExecutorService executor
+  ) {
+    this.isEnabled = isEnabled;
+    this.supplyMap = supplyMap;
+    this.zkValues = new AtomicReference<>(new HashMap<>());
+
+    if (this.isEnabled) {
+      executor.scheduleAtFixedRate(this::reloadZkValues, 2, cacheTtl, TimeUnit.SECONDS);
+    }
+  }
+
+  private void reloadZkValues() {
+    try {
+      Map<K, V> newZkValues = supplyMap.get();
+      zkValues.set(newZkValues);
+    } catch (Exception e) {
+      LOG.warn("Reloading ApiCache failed: {}", e.getMessage());
+    }
+  }
+
+  public V get(K key) {
+    return this.zkValues.get().get(key);
+  }
+
+  public Map<K, V> getAll() {
+    Map<K, V> allValues = this.zkValues.get();
+    if (allValues.isEmpty()) {
+      LOG.debug("ApiCache getAll returned empty");
+    }
+    return allValues;
+  }
+
+  public Map<K, V> getAll(Collection<K> keys) {
+    Map<K, V> allValues = this.zkValues.get();
+    Map<K, V> filteredValues = keys
+      .stream()
+      .filter(allValues::containsKey)
+      .collect(Collectors.toMap(Function.identity(), allValues::get));
+
+    if (filteredValues.isEmpty()) {
+      LOG.debug("ApiCache getAll returned empty for {}", keys);
+    }
+
+    return filteredValues;
+  }
+
+  public boolean isEnabled() {
+    return isEnabled;
+  }
+}

SingularityService/src/main/java/com/hubspot/singularity/data/DeployManager.java

@@ -17,6 +17,7 @@
 import com.hubspot.singularity.SingularityDeployStatistics;
 import com.hubspot.singularity.SingularityDeployUpdate;
 import com.hubspot.singularity.SingularityDeployUpdate.DeployEventType;
+import com.hubspot.singularity.SingularityManagedScheduledExecutorServiceFactory;
 import com.hubspot.singularity.SingularityPendingDeploy;
 import com.hubspot.singularity.SingularityRequest;
 import com.hubspot.singularity.SingularityRequestDeployState;
@@ -70,6 +71,8 @@ public class DeployManager extends CuratorAsyncManager {
   private static final String DEPLOY_STATISTICS_KEY = "STATISTICS";
   private static final String DEPLOY_RESULT_KEY = "RESULT_STATE";
 
+  private final ApiCache<String, SingularityRequestDeployState> deployCache;
+
   @Inject
   public DeployManager(
     CuratorFramework curator,
@@ -85,7 +88,8 @@ public DeployManager(
     IdTranscoder<SingularityDeployKey> deployKeyTranscoder,
     Transcoder<SingularityUpdatePendingDeployRequest> updateRequestTranscoder,
     ZkCache<SingularityDeploy> deploysCache,
-    SingularityLeaderCache leaderCache
+    SingularityLeaderCache leaderCache,
+    SingularityManagedScheduledExecutorServiceFactory executorServiceFactory
   ) {
     super(curator, configuration, metricRegistry);
     this.singularityEventListener = singularityEventListener;
@@ -99,6 +103,13 @@ public DeployManager(
     this.updateRequestTranscoder = updateRequestTranscoder;
     this.deploysCache = deploysCache;
     this.leaderCache = leaderCache;
+    this.deployCache =
+      new ApiCache<>(
+        configuration.useApiCacheInDeployManager(),
+        configuration.getDeployCacheTtl(),
+        this::fetchAllDeployStates,
+        executorServiceFactory.get("deploy-api-cache-reloader")
+      );
   }
 
   public List<SingularityDeployKey> getDeployIdsFor(String requestId) {
@@ -116,6 +127,11 @@ public List<SingularityDeployKey> getAllDeployIds() {
     return getChildrenAsIdsForParents("getAllDeployIds", paths, deployKeyTranscoder);
   }
 
+  public Map<String, SingularityRequestDeployState> fetchAllDeployStates() {
+    final List<String> requestIds = getChildren(BY_REQUEST_ROOT);
+    return fetchDeployStatesByRequestIds(requestIds);
+  }
+
   @Timed
   public Map<String, SingularityRequestDeployState> getRequestDeployStatesByRequestIds(
     Collection<String> requestIds
@@ -124,6 +140,15 @@ public Map<String, SingularityRequestDeployState> getRequestDeployStatesByReques
       return leaderCache.getRequestDeployStateByRequestId(requestIds);
     }
 
+    Map<String, SingularityRequestDeployState> deployStatesByRequestIds;
+
+    if (deployCache.isEnabled()) {
+      deployStatesByRequestIds = deployCache.getAll(requestIds);
+      if (!deployStatesByRequestIds.isEmpty()) {
+        return deployStatesByRequestIds;
+      }
+    }
+
     return fetchDeployStatesByRequestIds(requestIds);
   }
 

SingularityService/src/main/java/com/hubspot/singularity/data/RequestManager.java

@@ -13,6 +13,7 @@
 import com.hubspot.singularity.SingularityCreateResult;
 import com.hubspot.singularity.SingularityDeleteResult;
 import com.hubspot.singularity.SingularityDeployKey;
+import com.hubspot.singularity.SingularityManagedScheduledExecutorServiceFactory;
 import com.hubspot.singularity.SingularityPendingRequest;
 import com.hubspot.singularity.SingularityPendingRequest.PendingType;
 import com.hubspot.singularity.SingularityRequest;
@@ -32,11 +33,13 @@
 import com.hubspot.singularity.expiring.SingularityExpiringScale;
 import com.hubspot.singularity.expiring.SingularityExpiringSkipHealthchecks;
 import com.hubspot.singularity.scheduler.SingularityLeaderCache;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.utils.ZKPaths;
@@ -90,6 +93,7 @@ public class RequestManager extends CuratorAsyncManager {
   );
 
   private final Map<Class<? extends SingularityExpiringRequestActionParent<? extends SingularityExpiringRequestParent>>, Transcoder<? extends SingularityExpiringRequestActionParent<? extends SingularityExpiringRequestParent>>> expiringTranscoderMap;
+  private final ApiCache<String, SingularityRequestWithState> requestsCache;
 
   @Inject
   public RequestManager(
@@ -108,7 +112,8 @@ public RequestManager(
     Transcoder<SingularityExpiringSkipHealthchecks> expiringSkipHealthchecksTranscoder,
     SingularityWebCache webCache,
     SingularityLeaderCache leaderCache,
-    Transcoder<CrashLoopInfo> crashLoopInfoTranscoder
+    Transcoder<CrashLoopInfo> crashLoopInfoTranscoder,
+    SingularityManagedScheduledExecutorServiceFactory scheduledExecutorServiceFactory
   ) {
     super(curator, configuration, metricRegistry);
     this.requestTranscoder = requestTranscoder;
@@ -133,6 +138,16 @@ public RequestManager(
 
     this.leaderCache = leaderCache;
     this.webCache = webCache;
+    this.requestsCache =
+      new ApiCache<>(
+        configuration.useApiCacheInRequestManager(),
+        configuration.getRequestCacheTtl(),
+        () ->
+          fetchRequests()
+            .stream()
+            .collect(Collectors.toMap(r -> r.getRequest().getId(), Function.identity())),
+        scheduledExecutorServiceFactory.get("request-api-cache-reloader")
+      );
   }
 
   private String getRequestPath(String requestId) {
@@ -632,11 +647,22 @@ public List<SingularityRequestWithState> getRequests(boolean useWebCache) {
     if (useWebCache && webCache.useCachedRequests()) {
       return webCache.getRequests();
     }
+
+    if (requestsCache.isEnabled()) {
+      List<SingularityRequestWithState> requests = new ArrayList<>(
+        (requestsCache.getAll()).values()
+      );
+      if (!requests.isEmpty()) {
+        return requests;
+      }
+    }
+
     List<SingularityRequestWithState> requests = fetchRequests();
 
     if (useWebCache) {
       webCache.cacheRequests(requests);
     }
+
     return requests;
   }