Kafka AMI promotion workflow by Orion

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

@@ -0,0 +1,100 @@
+package com.pinterest.orion.core.actions.aws;
+
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import com.pinterest.orion.server.api.AMI;
+
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.services.ec2.Ec2Client;
+import software.amazon.awssdk.services.ec2.model.DescribeImagesRequest;
+import software.amazon.awssdk.services.ec2.model.DescribeImagesResponse;
+import software.amazon.awssdk.services.ec2.model.CreateTagsRequest;
+import software.amazon.awssdk.services.ec2.model.CreateTagsResponse;
+import software.amazon.awssdk.services.ec2.model.Filter;
+import software.amazon.awssdk.services.ec2.model.Tag;
+
+public class AmiTagManager {
+  private static final Logger logger = Logger.getLogger(AmiTagManager.class.getCanonicalName());
+  private Ec2Client ec2Client;
+
+  public AmiTagManager() {
+    ec2Client = Ec2Client.create();
+  }
+
+  public List<AMI> getAmiList(Map<String, String> filter) {
+    List<AMI> amiList = new ArrayList<>();
+    DescribeImagesRequest.Builder builder = DescribeImagesRequest.builder();
+    builder = builder.filters(
+      Filter.builder().name("tag:application").values("kafka").build()
+    );
+    if (filter.containsKey("os"))
+      builder = builder.filters(
+        Filter.builder().name("tag:release").values(filter.get("os")).build()
+      );
+    if (filter.containsKey("arch"))
+      builder = builder.filters(
+        Filter.builder().name("tag:cpu_architecture").values(filter.get("arch")).build()
+      );
+    builder = builder.filters(
+      Filter.builder().name("tag:application_environment").values("*").build()
+    );
+    try {
+      DescribeImagesResponse resp = ec2Client.describeImages(builder.build());
+      if (resp.hasImages() && !resp.images().isEmpty()) {
+        ZonedDateTime cutDate = ZonedDateTime.now().minusDays(180);
+        resp.images().forEach(image -> {
+          if (ZonedDateTime.parse(image.creationDate(), DateTimeFormatter.ISO_ZONED_DATE_TIME).isAfter(cutDate)) {
+            Iterator<Tag> i = image.tags().iterator();
+            Tag t;
+            String appEnvTag = null;
+            while (i.hasNext()) {
+              t = i.next();
+              if (t.key().equals("application_environment")) {
+                appEnvTag = t.value();
+                break;
+              }
+            }
+            amiList.add(new AMI(
+              image.imageId(),
+              appEnvTag,
+              image.creationDate()
+            ));
+          }
+        });
+        amiList.sort((a, b) -> - ZonedDateTime.parse(a.getCreationDate(), DateTimeFormatter.ISO_ZONED_DATE_TIME)
+        .compareTo(ZonedDateTime.parse(b.getCreationDate(), DateTimeFormatter.ISO_ZONED_DATE_TIME)));
+      }
+    } catch (Exception e) {
+      logger.log(Level.SEVERE, "AMITagManager: could not retrieve AMI list", e);
+      throw e;
+    }
+    return amiList;
+  }
+
+  public void updateAmiTag(String amiId, String applicationEnvironment) {
+    CreateTagsRequest.Builder builder = CreateTagsRequest.builder();
+    Tag newAppEnv = Tag.builder()
+        .key("application_environment")
+        .value(applicationEnvironment).build();
+    CreateTagsRequest request = builder
+        .resources(amiId)
+        .tags(newAppEnv)
+        .build();
+    CreateTagsResponse resp;
+    try {
+      resp = ec2Client.createTags(request);
+      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);
+      throw e;
+    }
+  }
+}

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

@@ -0,0 +1,25 @@
+package com.pinterest.orion.server.api;
+
+public class AMI {
+  private final String amiId;
+  private final String applicationEnvironment;
+  private final String creationDate;
+
+  public AMI(String amiId, String applicationEnvironment, String creationDate) {
+    this.amiId = amiId;
+    this.applicationEnvironment = applicationEnvironment;
+    this.creationDate = creationDate;
+  }
+
+  public String getAmiId() {
+    return amiId;
+  }
+
+  public String getApplicationEnvironment() {
+    return applicationEnvironment;
+  }
+
+  public String getCreationDate() {
+    return creationDate;
+  }
+}

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

@@ -25,9 +25,11 @@
 
 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;
@@ -41,11 +43,13 @@
 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);
@@ -111,6 +115,33 @@ public Map<String, Map<String, Utilization>> getUtilizationDetailsByCluster() {
     return utilizationMap;
   }
 
+  @Path("/describeImages")
+  @GET
+  public List<AMI> describeImages(
+      @QueryParam("os") String os,
+      @QueryParam("arch") String arch
+  ) {
+    Map<String, String> filter = new HashMap<>();
+    if (os != null)
+      filter.put("os", os);
+    if (arch != null)
+      filter.put("arch", arch);
+    if (amiTagManager == null)
+      amiTagManager = new AmiTagManager();
+    return amiTagManager.getAmiList(filter);
+  }
+
+  @Path("/updateImageTag")
+  @PUT
+  public void updateImageTag(
+      @QueryParam("ami_id") String amiId,
+      @QueryParam("application_environment") String applicationEnvironment
+  ) {
+    if (amiTagManager == null)
+      amiTagManager = new AmiTagManager();
+    amiTagManager.updateAmiTag(amiId, applicationEnvironment);
+  }
+
   @RolesAllowed({ OrionConf.ADMIN_ROLE, OrionConf.MGMT_ROLE })
   @Path("/costByCluster")
   @GET

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

@@ -30,6 +30,9 @@ 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 AMILIST_REQUESTED = "AMILIST_REQUESTED";
+export const AMILIST_RECEIVED = "AMILIST_RECEIVED";
+export const AMITAG_UPDATE = "AMITAG_UPDATE";
 
 export function requestCluster(clusterId) {
   return { type: CLUSTER_REQUESTED, payload: { clusterId } };
@@ -110,3 +113,24 @@ export function receiveClusterEndpoint(clusterId, field, data) {
     payload: { clusterId, field, data },
   };
 }
+
+export function requestAmiList(filter) {
+  return {
+    type: AMILIST_REQUESTED,
+    payload: { filter },
+  };
+}
+
+export function receiveAmiList(amiList) {
+  return {
+    type: AMILIST_RECEIVED,
+    payload: { amiList },
+  };
+}
+
+export function updateAmiTag(amiId, applicationEnvironment) {
+  return {
+    type: AMITAG_UPDATE,
+    payload: { amiId, applicationEnvironment },
+  };
+}