Kafka AMI promotion workflow by Orion

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

@@ -0,0 +1,141 @@
+/*******************************************************************************
+ * Copyright 2024 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.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 java.util.function.UnaryOperator;
+
+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;
+
+/**
+ * AmiTagManager interfaces APIs in
+ * {@link com.pinterest.orion.server.api.ClusterManagerApi ClusterManagerApi}
+ * with AWS for AMI tag management.
+ */
+public class AmiTagManager {
+  private static final Logger logger = Logger.getLogger(AmiTagManager.class.getCanonicalName());
+  private Ec2Client ec2Client;
+  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";
+  public static final String KEY_RELEASE = "release";
+  public static final String KEY_CPU_ARCHITECTURE = "cpu_architecture";
+  public static final String KEY_APPLICATION_ENVIRONMENT = "application_environment";
+  public static final String VALUE_KAFKA = "kafka";
+  public static UnaryOperator<String> tag = key -> "tag:" + key;
+
+  public AmiTagManager() {
+    ec2Client = Ec2Client.create();
+  }
+
+  /**
+   * 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) {
+    List<Ami> amiList = new ArrayList<>();
+    DescribeImagesRequest.Builder builder = DescribeImagesRequest.builder();
+    builder = builder.filters(
+      Filter.builder().name(tag.apply(KEY_APPLICATION)).values(VALUE_KAFKA).build()
+    );
+    if (filter.containsKey(KEY_RELEASE))
+      builder = builder.filters(
+        Filter.builder().name(tag.apply(KEY_RELEASE)).values(filter.get(KEY_RELEASE)).build()
+      );
+    if (filter.containsKey(KEY_CPU_ARCHITECTURE))
+      builder = builder.filters(
+        Filter.builder().name(tag.apply(KEY_CPU_ARCHITECTURE)).values(filter.get(KEY_CPU_ARCHITECTURE)).build()
+      );
+    builder = builder.filters(
+      Filter.builder().name(tag.apply(KEY_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(KEY_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;
+  }
+
+  /**
+   * update AMI tag 'application_environment'
+   *
+   * @param amiId - target AMI id
+   * @param applicationEnvironment - new tag value
+   */
+  public void updateAmiTag(String amiId, String applicationEnvironment) {
+    CreateTagsRequest.Builder builder = CreateTagsRequest.builder();
+    Tag newAppEnv = Tag.builder()
+        .key(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,53 @@
+/*******************************************************************************
+ * Copyright 2024 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;
+
+/**
+ * Specialized object to store AMI information.
+ * Used to transmit a list of AMIs to the frontend.
+ */
+public class Ami {
+  private final String amiId;
+  private final String applicationEnvironment;
+  private final String creationDate;
+
+  /**
+   * creates a new Ami instance
+   *
+   * @param amiId - for AWS, the string "ami-" and a sequence of 17 characters
+   * @param applicationEnvironment - comma-separated list of environments (dev, test,
+   *                                 staging, prod) supported by this ami
+   * @param creationDate - ami creation date, UTC
+   * @return the new Ami instance
+   */
+  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(AmiTagManager.KEY_RELEASE) String os,
+      @QueryParam(AmiTagManager.KEY_CPU_ARCHITECTURE) String arch
+  ) {
+    Map<String, String> filter = new HashMap<>();
+    if (os != null)
+      filter.put(AmiTagManager.KEY_RELEASE, os);
+    if (arch != null)
+      filter.put(AmiTagManager.KEY_CPU_ARCHITECTURE, arch);
+    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);
+  }
+
   @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 AMI_LIST_REQUESTED = "AMI_LIST_REQUESTED";
+export const AMI_LIST_RECEIVED = "AMI_LIST_RECEIVED";
+export const AMI_TAG_UPDATE = "AMI_TAG_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: 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 },
+  };
+}