Get cluster capacity in Teletraan backend from Rodimus API

pinterest · Jul 30, 2024 · bf3f469 · bf3f469
1 parent 5fc922d
commit bf3f469
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 1 deletion.
diff --git a/deploy-service/common/src/main/java/com/pinterest/deployservice/handler/CommonHandler.java b/deploy-service/common/src/main/java/com/pinterest/deployservice/handler/CommonHandler.java
@@ -23,6 +23,7 @@
 import com.pinterest.deployservice.common.WebhookDataFactory;
 import com.pinterest.deployservice.dao.*;
 import com.pinterest.deployservice.email.MailManager;
+import com.pinterest.deployservice.rodimus.RodimusManager;
 import com.pinterest.deployservice.events.DeployEvent;
 import com.pinterest.teletraan.universal.events.AppEventPublisher;
 
@@ -60,6 +61,8 @@ public class CommonHandler {
     private Counter successCounter;
     private Counter failureCounter;
     private BuildTagsManager buildTagsManager;
+    private final RodimusManager rodimusManager;
+
 
     private final class FinishNotifyJob implements Callable<Void> {
         private EnvironBean envBean;
@@ -130,6 +133,7 @@ public CommonHandler(ServiceContext serviceContext) {
         chatManager = serviceContext.getChatManager();
         mailManager = serviceContext.getMailManager();
         buildTagsManager = serviceContext.getBuildTagsManager();
+        rodimusManager = serviceContext.getRodimusManager();
         jobPool = serviceContext.getJobPool();
         dataHandler = new DataHandler(serviceContext);
         deployBoardUrlPrefix = serviceContext.getDeployBoardUrlPrefix();
@@ -311,8 +315,12 @@ void transition(DeployBean deployBean, DeployBean newDeployBean, EnvironBean env
         newDeployBean.setState(oldState);
         newDeployBean.setLast_update(System.currentTimeMillis());
 
+        String cluster = envBean.getEnv_name() + "-" + envBean.getStage_name();
+        Long capacity = rodimusManager.getClusterCapacity(cluster);
+        LOG.debug("Yaqin Debug: The capacity for cluster {} is {}", cluster, capacity);
+
         //The maximum sucThreshold is 10000 to keep precision.
-        if (succeeded * 10000 >= sucThreshold * total && succeeded > 0) {
+        if (succeeded * 10000 >= sucThreshold * total && !(succeeded == 0 && capacity != null && capacity > 0)) {
             LOG.debug("Propose deploy {} as SUCCEEDING since {} agents are succeeded.", deployId, succeeded);
             if (deployBean.getSuc_date() == null) {
                 newDeployBean.setSuc_date(System.currentTimeMillis());

diff --git a/...rvice/common/src/main/java/com/pinterest/deployservice/rodimus/DefaultRodimusManager.java b/...rvice/common/src/main/java/com/pinterest/deployservice/rodimus/DefaultRodimusManager.java
@@ -47,4 +47,9 @@ public Long getClusterInstanceLaunchGracePeriod(String clusterName) throws Excep
     public Map<String, Map<String, String>> getEc2Tags(Collection<String> hostIds) throws Exception {
         return new HashMap<>();
     }
+
+    @Override
+    public Long getClusterCapacity(String clusterName) throws Exception {
+        return null;
+    }
 }
diff --git a/deploy-service/common/src/main/java/com/pinterest/deployservice/rodimus/RodimusManager.java b/deploy-service/common/src/main/java/com/pinterest/deployservice/rodimus/RodimusManager.java
@@ -31,4 +31,6 @@ void terminateHostsByClusterName(String clusterName, Collection<String> hostIds,
     Long getClusterInstanceLaunchGracePeriod(String clusterName) throws Exception;
 
     Map<String, Map<String, String>> getEc2Tags(Collection<String> hostIds) throws Exception;
+
+    Long getClusterCapacity(String clusterName) throws Exception;
 }
diff --git a/...-service/common/src/main/java/com/pinterest/deployservice/rodimus/RodimusManagerImpl.java b/...-service/common/src/main/java/com/pinterest/deployservice/rodimus/RodimusManagerImpl.java
@@ -179,4 +179,27 @@ public Map<String, Map<String, String>> getEc2Tags(Collection<String> hostIds) t
         }.getType());
     } // getEc2Tags
 
+    @Override
+    public Long getClusterCapacity(String clusterName) throws Exception {
+        String url = String.format("%s/v1/groups/%s", rodimusUrl, clusterName);
+        String res = callHttpClient(Verb.GET, url, null);
+
+        JsonObject jsonObject = gson.fromJson(res, JsonObject.class);
+        if (jsonObject == null || jsonObject.isJsonNull()) {
+            return null;
+        }
+
+        JsonObject launchInfo = jsonObject.getAsJsonObject("launchInfo");
+        if (launchInfo == null || launchInfo.isJsonNull()) {
+            return null;
+        }
+
+        JsonPrimitive capacity = launchInfo.getAsJsonPrimitive("minSize");
+        if (capacity == null || capacity.isJsonNull()) {
+            return null;
+        }
+
+        return capacity.getAsLong();
+    } // getClusterCapacity
+
 } // class RodimusManagerImpl