Fix AMI query timeout

orion-server/src/main/java/com/pinterest/orion/core/actions/aws/AmiTagManager.java

@@ -21,6 +21,10 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.function.Function;
@@ -45,6 +49,8 @@
 public class AmiTagManager {
   private static final Logger logger = Logger.getLogger(AmiTagManager.class.getCanonicalName());
   private Ec2Client ec2Client;
+  private ExecutorService executorService = Executors.newCachedThreadPool();
+  private Map<String, CompletableFuture<List<Ami>>> activeFutures = new ConcurrentHashMap<>();
   public static final String KEY_IS_PUBLIC = "is-public";
   public static final String KEY_AMI_ID = "ami_id";
   public static final String KEY_APPLICATION = "application";
@@ -60,13 +66,37 @@ public AmiTagManager() {
     ec2Client = Ec2Client.create();
   }
 
+  /**
+   * run queries in a pool of threads
+   *
+   * @param filter - map of criteria fields
+   * @return AMI list, or null if waiting to complete the query
+   */
+  public List<Ami> getAmiListAsync(Map<String, String> filter) throws Exception {
+    String filterKey = filter.toString();
+    CompletableFuture<List<Ami>> future = activeFutures.get(filterKey);
+    // Start a new query if it doesn't already exist
+    if (future == null) {
+      future = CompletableFuture.supplyAsync(() -> {
+        return getAmiList(filter);
+      }, executorService);
+      activeFutures.put(filterKey, future);
+    }
+
+    if (future.isDone()) {
+      activeFutures.remove(filterKey);
+      return future.get();
+    }
+    return null;
+  }
+
   /**
    * retrieve AMI list from cloud provider
    *
    * @param filter - map of criteria fields
    * @return list of Ami objects
    */
-  public List<Ami> getAmiList(Map<String, String> filter) {
+  private List<Ami> getAmiList(Map<String, String> filter) {
     List<Ami> amiList = new ArrayList<>();
     DescribeImagesRequest.Builder builder = DescribeImagesRequest.builder();
     Filter.Builder filterBuilder = Filter.builder();
@@ -101,6 +131,7 @@ public List<Ami> getAmiList(Map<String, String> filter) {
     );
     builder = builder.filters(filterList);
     try {
+      logger.info("Querying AWS for AMI list");
       DescribeImagesResponse resp = ec2Client.describeImages(builder.build());
       if (resp.hasImages() && !resp.images().isEmpty()) {
         // The limitation of images newer than 180 days is temporarily suspended
@@ -127,8 +158,9 @@ public List<Ami> getAmiList(Map<String, String> filter) {
         Function<Ami, ZonedDateTime> parse = i -> ZonedDateTime.parse(i.getCreationDate(), DateTimeFormatter.ISO_ZONED_DATE_TIME);
         amiList.sort((a, b) -> - parse.apply(a).compareTo(parse.apply(b)));
       }
+      logger.info("AWS query complete");
     } catch (Exception e) {
-      logger.log(Level.SEVERE, "AmiTagManager: could not retrieve AMI list", e);
+      logger.log(Level.SEVERE, "Could not retrieve AMI list", e);
       throw e;
     }
     return amiList;
@@ -155,7 +187,7 @@ public void updateAmiTag(String amiId, String applicationEnvironment) {
       if (!resp.sdkHttpResponse().isSuccessful())
         throw AwsServiceException.builder().message("Http code \" + resp.sdkHttpResponse().statusCode() + \" received").build();
     } catch (Exception e) {
-      logger.severe("AmiTagManager: tag update failed for " + amiId + " and application_environment tag = " + applicationEnvironment + ", " + e);
+      logger.severe("Tag update failed for " + amiId + " and application_environment tag = " + applicationEnvironment + ", " + e);
       throw e;
     }
   }

orion-server/src/main/java/com/pinterest/orion/server/OrionServer.java

@@ -51,6 +51,7 @@
 import com.pinterest.orion.security.NoopAuthorizationFilter;
 import com.pinterest.orion.security.OrionAuthorizationFilter;
 import com.pinterest.orion.server.api.ActionEngineApi;
+import com.pinterest.orion.server.api.AmiApi;
 import com.pinterest.orion.server.api.ClusterApi;
 import com.pinterest.orion.server.api.ClusterManagerApi;
 import com.pinterest.orion.server.api.CustomApiFactory;
@@ -275,6 +276,7 @@ private void registerAPIs(Environment environment, OrionConf configuration) {
                                                                  // the future and the admin api
                                                                  // should be the endpoint for
                                                                  // registeration
+    environment.jersey().register(new AmiApi(mgr));
     if (configuration.getCustomApiFactoryClasses() != null) {
       for (String factoryClass : configuration.getCustomApiFactoryClasses()) {
         CustomApiFactory instance;

orion-server/src/main/java/com/pinterest/orion/server/api/AmiApi.java

@@ -0,0 +1,88 @@
+/*******************************************************************************
+ * Copyright 2020 Pinterest, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *******************************************************************************/
+package com.pinterest.orion.server.api;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import javax.ws.rs.GET;
+import javax.ws.rs.PUT;
+import javax.ws.rs.Path;
+import javax.ws.rs.Produces;
+import javax.ws.rs.QueryParam;
+import javax.ws.rs.core.MediaType;
+import com.pinterest.orion.core.ClusterManager;
+import com.pinterest.orion.core.actions.aws.AmiTagManager;
+
+@Path("/")
+@Produces({ MediaType.APPLICATION_JSON })
+public class AmiApi extends BaseClustersApi {
+  private AmiTagManager amiTagManager = new AmiTagManager();
+
+  public AmiApi(ClusterManager mgr) {
+    super(mgr);
+  }
+
+  @Path("/describeImages")
+  @GET
+  /**
+   * @param os: bionic, focal, release
+   * @param arch: x86_64, arm64
+   * @param environment: dev, test, stage, prod
+   * @return AMI list, or null if waiting for query completion
+   * @throws Exception
+   * 
+   * Http response status:
+   * - AMI list: 200
+   * - null: 204
+   */
+  public List<Ami> describeImages(
+      @QueryParam(AmiTagManager.KEY_RELEASE) String os,
+      @QueryParam(AmiTagManager.KEY_ARCHITECTURE) String arch,
+      @QueryParam(AmiTagManager.KEY_ENVIRONMENT) String environment
+  ) throws Exception {
+    Map<String, String> filter = new LinkedHashMap<>();
+    if (os != null)
+      filter.put(AmiTagManager.KEY_RELEASE, os);
+    if (arch != null)
+      filter.put(AmiTagManager.KEY_ARCHITECTURE, arch);
+    if (environment != null)
+      filter.put(AmiTagManager.KEY_ENVIRONMENT, environment);
+    return amiTagManager.getAmiListAsync(filter);
+  }
+
+  @Path("/updateImageTag")
+  @PUT
+  public void updateImageTag(
+      @QueryParam(AmiTagManager.KEY_AMI_ID) String amiId,
+      @QueryParam(AmiTagManager.KEY_APPLICATION_ENVIRONMENT) String applicationEnvironment
+  ) {
+    amiTagManager.updateAmiTag(amiId, applicationEnvironment);
+  }
+
+  @Path("/getEnvTypes")
+  @GET
+  public List<String> getEnvTypes() {
+    List<String> envTypes = null;
+    Map<String, Object> additionalConfigs = mgr.getOrionConf().getAdditionalConfigs();
+    if(additionalConfigs != null && additionalConfigs.containsKey(AmiTagManager.ENV_TYPES_KEY)) {
+      envTypes = (List<String>) additionalConfigs.get(AmiTagManager.ENV_TYPES_KEY);
+    }
+    return envTypes;
+  }
+
+}

orion-server/src/main/java/com/pinterest/orion/server/api/ClusterManagerApi.java

@@ -25,11 +25,9 @@
 
 import javax.annotation.security.RolesAllowed;
 import javax.ws.rs.GET;
-import javax.ws.rs.PUT;
 import javax.ws.rs.NotFoundException;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
-import javax.ws.rs.QueryParam;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.SecurityContext;
@@ -43,13 +41,11 @@
 import com.pinterest.orion.core.Utilization;
 import com.pinterest.orion.core.global.sensor.GlobalPluginManager;
 import com.pinterest.orion.core.global.sensor.GlobalSensor;
-import com.pinterest.orion.core.actions.aws.AmiTagManager;
 import com.pinterest.orion.server.config.OrionConf;
 
 @Path("/")
 @Produces({ MediaType.APPLICATION_JSON })
 public class ClusterManagerApi extends BaseClustersApi {
-  private AmiTagManager amiTagManager;
 
   public ClusterManagerApi(ClusterManager mgr) {
     super(mgr);
@@ -115,48 +111,6 @@ public Map<String, Map<String, Utilization>> getUtilizationDetailsByCluster() {
     return utilizationMap;
   }
 
-  @Path("/describeImages")
-  @GET
-  public List<Ami> describeImages(
-      @QueryParam(AmiTagManager.KEY_RELEASE) String os,
-      @QueryParam(AmiTagManager.KEY_ARCHITECTURE) String arch,
-      @QueryParam(AmiTagManager.KEY_ENVIRONMENT) String environment
-  ) {
-    Map<String, String> filter = new HashMap<>();
-    if (os != null)
-      filter.put(AmiTagManager.KEY_RELEASE, os);
-    if (arch != null)
-      filter.put(AmiTagManager.KEY_ARCHITECTURE, arch);
-    if (environment != null)
-      filter.put(AmiTagManager.KEY_ENVIRONMENT, environment);
-    if (amiTagManager == null)
-      amiTagManager = new AmiTagManager();
-    return amiTagManager.getAmiList(filter);
-  }
-
-  @Path("/updateImageTag")
-  @PUT
-  public void updateImageTag(
-      @QueryParam(AmiTagManager.KEY_AMI_ID) String amiId,
-      @QueryParam(AmiTagManager.KEY_APPLICATION_ENVIRONMENT) String applicationEnvironment
-  ) {
-    if (amiTagManager == null)
-      amiTagManager = new AmiTagManager();
-    amiTagManager.updateAmiTag(amiId, applicationEnvironment);
-  }
-
-
-  @Path("/getEnvTypes")
-  @GET
-  public List<String> getEnvTypes() {
-    List<String> envTypes = null;
-    Map<String, Object> additionalConfigs = mgr.getOrionConf().getAdditionalConfigs();
-    if(additionalConfigs != null && additionalConfigs.containsKey(AmiTagManager.ENV_TYPES_KEY)) {
-      envTypes = (List<String>) additionalConfigs.get(AmiTagManager.ENV_TYPES_KEY);
-    }
-    return envTypes;
-  }
-
   @RolesAllowed({ OrionConf.ADMIN_ROLE, OrionConf.MGMT_ROLE })
   @Path("/costByCluster")
   @GET

orion-server/src/main/resources/webapp/src/actions/ami.js

@@ -0,0 +1,55 @@
+/*******************************************************************************
+ * Copyright 2020 Pinterest, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *******************************************************************************/
+export const AMI_LIST_REQUESTED = "AMI_LIST_REQUESTED";
+export const AMI_LIST_RECEIVED = "AMI_LIST_RECEIVED";
+export const AMI_TAG_UPDATE = "AMI_TAG_UPDATE";
+export const ENV_TYPES_REQUESTED = "ENV_TYPES_REQUESTED";
+export const ENV_TYPES_RECEIVED = "ENV_TYPES_RECEIVED";
+
+export function requestAmiList(filter) {
+  return {
+    type: AMI_LIST_REQUESTED,
+    payload: { filter },
+  };
+}
+
+export function receiveAmiList(amiList) {
+  return {
+    type: AMI_LIST_RECEIVED,
+    payload: { amiList },
+  };
+}
+
+export function updateAmiTag(amiList, amiId, applicationEnvironment) {
+  return {
+    type: AMI_TAG_UPDATE,
+    payload: { amiList, amiId, applicationEnvironment },
+  };
+}
+
+export function requestEnvTypes() {
+  return {
+    type: ENV_TYPES_REQUESTED,
+    payload: {},
+  };
+}
+
+export function receiveEnvTypes(envTypeList) {
+  return {
+    type: ENV_TYPES_RECEIVED,
+    payload: { envTypeList },
+  };
+}

orion-server/src/main/resources/webapp/src/actions/cluster.js

@@ -30,11 +30,6 @@ export const UTILIZATION_REQUESTED = "UTILIZATION_REQUESTED";
 export const UTILIZATION_RECEIVED = "UTILIZATION_RECEIVED";
 export const COST_REQUESTED = "COST_REQUESTED";
 export const COST_RECEIVED = "COST_RECEIVED";
-export const AMI_LIST_REQUESTED = "AMI_LIST_REQUESTED";
-export const AMI_LIST_RECEIVED = "AMI_LIST_RECEIVED";
-export const AMI_TAG_UPDATE = "AMI_TAG_UPDATE";
-export const ENV_TYPES_REQUESTED = "ENV_TYPES_REQUESTED";
-export const ENV_TYPES_RECEIVED = "ENV_TYPES_RECEIVED";
 
 export function requestCluster(clusterId) {
   return { type: CLUSTER_REQUESTED, payload: { clusterId } };
@@ -115,38 +110,3 @@ export function receiveClusterEndpoint(clusterId, field, data) {
     payload: { clusterId, field, data },
   };
 }
-
-export function requestAmiList(filter) {
-  return {
-    type: AMI_LIST_REQUESTED,
-    payload: { filter },
-  };
-}
-
-export function receiveAmiList(amiList) {
-  return {
-    type: AMI_LIST_RECEIVED,
-    payload: { amiList },
-  };
-}
-
-export function updateAmiTag(amiId, applicationEnvironment) {
-  return {
-    type: AMI_TAG_UPDATE,
-    payload: { amiId, applicationEnvironment },
-  };
-}
-
-export function requestEnvTypes() {
-  return {
-    type: ENV_TYPES_REQUESTED,
-    payload: {},
-  };
-}
-
-export function receiveEnvTypes(envTypeList) {
-  return {
-    type: ENV_TYPES_RECEIVED,
-    payload: { envTypeList },
-  };
-}