Merge pull request #159 from snallami/master

[JENKINS-27601] fix for instance caps which are calculated solely based on AMI ID
jenkinsci · Aug 18, 2015 · b27bea5 · b27bea5
2 parents 85e9152 + 18b8289
commit b27bea5
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 16 deletions.
diff --git a/src/main/java/hudson/plugins/ec2/EC2Cloud.java b/src/main/java/hudson/plugins/ec2/EC2Cloud.java
@@ -99,6 +99,8 @@ public abstract class EC2Cloud extends Cloud {
 
     public static final String DEFAULT_EC2_HOST = "us-east-1";
     public static final String EC2_URL_HOST = "ec2.amazonaws.com";
+    public static final String EC2_SLAVE_TYPE_SPOT = "spot";
+    public static final String EC2_SLAVE_TYPE_DEMAND = "demand";
 
     private final boolean useInstanceProfileForCredentials;
     private final String accessId;
@@ -218,18 +220,18 @@ public synchronized KeyPair getKeyPair() throws AmazonClientException, IOExcepti
     }
 
     /**
-     * Counts the number of instances in EC2 currently running that are using the specifed image.
+     * Counts the number of instances in EC2 currently running that are using the specifed image and a template.
      *
-     * @param ami
-     *            If AMI is left null, then all instances are counted.
-     *            <p>
-     *            This includes those instances that may be started outside Hudson.
+     * @param ami If AMI is left null, then all instances are counted and template description is ignored.
+     * <p>
+     * This includes those instances that may be started outside Hudson.
+     * @param templateDesc 
      */
-    public int countCurrentEC2Slaves(String ami) throws AmazonClientException {
-        int n = 0;
+    public int countCurrentEC2Slaves(String ami, String templateDesc) throws AmazonClientException {
+        int n=0;
         for (Reservation r : connect().describeInstances().getReservations()) {
             for (Instance i : r.getInstances()) {
-                if (isEc2ProvisionedSlave(i, ami)) {
+                if (isEc2ProvisionedAmiSlave(i, ami, templateDesc)) {
                     InstanceStateName stateName = InstanceStateName.fromValue(i.getState().getName());
                     if (stateName == InstanceStateName.Pending || stateName == InstanceStateName.Running) {
                         n++;
@@ -240,13 +242,23 @@ public int countCurrentEC2Slaves(String ami) throws AmazonClientException {
         return n;
     }
 
-    protected boolean isEc2ProvisionedSlave(Instance i, String ami) {
+    protected boolean isEc2ProvisionedAmiSlave(Instance i, String ami, String templateDesc) {
         // Check if the ami matches
         if (ami == null || StringUtils.equals(ami, i.getImageId())) {
             // Check if there is a ec2slave tag...
             for (Tag tag : i.getTags()) {
                 if (StringUtils.equals(tag.getKey(), EC2Tag.TAG_NAME_JENKINS_SLAVE_TYPE)) {
-                    return true;
+                	if (ami == null || templateDesc == null) {
+                		return true;
+                	} else if (StringUtils.equals(tag.getValue(), EC2Cloud.EC2_SLAVE_TYPE_DEMAND) || StringUtils.equals(tag.getValue(), EC2Cloud.EC2_SLAVE_TYPE_SPOT)) {
+                		// To handle cases where description is null and also upgrade cases for existing slave nodes.
+                		return true;
+                	} else if (StringUtils.equals(tag.getValue(), getSlaveTypeTagValue(EC2Cloud.EC2_SLAVE_TYPE_DEMAND, templateDesc)) || 
+                			StringUtils.equals(tag.getValue(), getSlaveTypeTagValue(EC2Cloud.EC2_SLAVE_TYPE_SPOT, templateDesc))) {
+                		return true;
+                	} else {
+                		return false;
+                	}
                 }
             }
             return false;
@@ -295,10 +307,14 @@ public HttpResponse doProvision(@QueryParameter String template) throws ServletE
     /**
      * Check for the count of EC2 slaves and determine if a new slave can be added. Takes into account both what Amazon
      * reports as well as an internal count of slaves currently being "provisioned".
+     * Check for the count of EC2 slaves and determine if a new slave can be added.
+     * Takes into account both what Amazon reports as well as an internal count
+     * of slaves currently being "provisioned".
+     * @param templateDesc 
      */
-    private boolean addProvisionedSlave(String ami, int amiCap) throws AmazonClientException {
-        int estimatedTotalSlaves = countCurrentEC2Slaves(null);
-        int estimatedAmiSlaves = countCurrentEC2Slaves(ami);
+    private boolean addProvisionedSlave(String ami, int amiCap, String templateDesc) throws AmazonClientException {
+        int estimatedTotalSlaves = countCurrentEC2Slaves(null, null);
+        int estimatedAmiSlaves = countCurrentEC2Slaves(ami, templateDesc);
 
         synchronized (provisioningAmis) {
             int currentProvisioning;
@@ -381,7 +397,7 @@ public Collection<PlannedNode> provision(Label label, int excessWorkload) {
 
             while (excessWorkload > 0) {
 
-                if (!addProvisionedSlave(t.ami, amiCap)) {
+                if (!addProvisionedSlave(t.ami, amiCap, t.description)) {
                     break;
                 }
 
@@ -436,6 +452,10 @@ private AWSCredentialsProvider createCredentialsProvider() {
     public static AWSCredentialsProvider createCredentialsProvider(final boolean useInstanceProfileForCredentials, final String accessId, final String secretKey) {
         return createCredentialsProvider(useInstanceProfileForCredentials, accessId.trim(), Secret.fromString(secretKey.trim()));
     }
+
+    public static String getSlaveTypeTagValue(String slaveType, String templateDescription) {
+    	return templateDescription != null ? slaveType+"_"+templateDescription : slaveType;
+    }
 
     public static AWSCredentialsProvider createCredentialsProvider(final boolean useInstanceProfileForCredentials, final String accessId, final Secret secretKey) {
 

diff --git a/src/main/java/hudson/plugins/ec2/SlaveTemplate.java b/src/main/java/hudson/plugins/ec2/SlaveTemplate.java
@@ -462,7 +462,8 @@ private EC2AbstractSlave provisionOndemand(TaskListener listener) throws AmazonC
                 if (inst_tags == null) {
                     inst_tags = new HashSet<Tag>();
                 }
-                inst_tags.add(new Tag(EC2Tag.TAG_NAME_JENKINS_SLAVE_TYPE, "demand"));
+                // Append template description as well to identify slaves provisioned per template
+                inst_tags.add(new Tag(EC2Tag.TAG_NAME_JENKINS_SLAVE_TYPE, EC2Cloud.getSlaveTypeTagValue(EC2Cloud.EC2_SLAVE_TYPE_DEMAND, description)));
             }
 
             DescribeInstancesRequest diRequest = new DescribeInstancesRequest();
@@ -724,7 +725,7 @@ private EC2AbstractSlave provisionSpot(TaskListener listener) throws AmazonClien
             }
             if (!hasCustomTypeTag) {
                 if (inst_tags != null)
-                    inst_tags.add(new Tag(EC2Tag.TAG_NAME_JENKINS_SLAVE_TYPE, "spot"));
+            	    inst_tags.add(new Tag(EC2Tag.TAG_NAME_JENKINS_SLAVE_TYPE, EC2Cloud.getSlaveTypeTagValue(EC2Cloud.EC2_SLAVE_TYPE_SPOT, description)));
             }
 
             if (StringUtils.isNotBlank(getIamInstanceProfile())) {