Add option for single-use agents in pool retention strategy (#491)

Co-authored-by: Tim Jacomb <timjacomb1@gmail.com>
Co-authored-by: Tim Jacomb <21194782+timja@users.noreply.github.com>

src/main/java/com/microsoft/azure/vmagent/AzureVMCloudPoolRetentionStrategy.java

@@ -5,27 +5,34 @@
 import edu.umd.cs.findbugs.annotations.CheckForNull;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import hudson.Extension;
+import hudson.model.Executor;
+import hudson.model.ExecutorListener;
 import hudson.model.Computer;
+import hudson.model.Queue;
 import hudson.model.Descriptor;
 import hudson.model.DescriptorVisibilityFilter;
 import hudson.slaves.Cloud;
+import hudson.slaves.AbstractCloudComputer;
 import hudson.slaves.RetentionStrategy;
 import jenkins.model.Jenkins;
 import org.kohsuke.stapler.DataBoundConstructor;
+import org.kohsuke.stapler.DataBoundSetter;
 
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-public class AzureVMCloudPoolRetentionStrategy extends AzureVMCloudBaseRetentionStrategy {
+public class AzureVMCloudPoolRetentionStrategy extends AzureVMCloudBaseRetentionStrategy implements ExecutorListener {
     private static final long serialVersionUID = 1577788691L;
 
     private final long retentionMillis;
 
     private final int poolSize;
 
+    private boolean singleUseAgents;
+
     private static final long IDLE_LIMIT_MILLIS = TimeUnit.MINUTES.toMillis(1);
 
     private static final Logger LOGGER = Logger.getLogger(AzureVMManagementServiceDelegate.class.getName());
@@ -161,6 +168,58 @@ public int getPoolSize() {
         return poolSize;
     }
 
+    @Override
+    public void taskAccepted(Executor executor, Queue.Task task) {
+
+    }
+
+    @Override
+    public void taskCompleted(Executor executor, Queue.Task task, long durationMS) {
+        done(executor);
+    }
+
+    @Override
+    public void taskCompletedWithProblems(Executor executor, Queue.Task task, long durationMS, Throwable problems) {
+        done(executor);
+    }
+
+    public void done(Executor executor) {
+        final AbstractCloudComputer<?> computer = (AbstractCloudComputer) executor.getOwner();
+        if (!(computer instanceof AzureVMComputer)) {
+            return;
+        }
+        done((AzureVMComputer) computer);
+    }
+    public void done(AzureVMComputer computer) {
+        if (!isSingleUseAgents()) {
+            return;
+        }
+
+        final AzureVMAgent agent = computer.getNode();
+        if (agent == null) {
+            return;
+        }
+
+        computer.setAcceptingTasks(false);
+        if (agent.isShutdownOnIdle()) {
+            LOGGER.log(Level.FINE, "Tagging VM to shutdown when idle: {0}",
+                    computer.getName());
+            agent.setCleanUpAction(CleanUpAction.SHUTDOWN, Messages._Build_Action_Shutdown_Agent());
+        } else {
+            LOGGER.log(Level.FINE, "Tagging VM to delete when idle: {0}", computer.getName());
+            agent.setCleanUpAction(CleanUpAction.DELETE, Messages._Build_Action_Delete_Agent());
+        }
+    }
+
+    public boolean isSingleUseAgents() {
+        return this.singleUseAgents;
+    }
+
+    @DataBoundSetter
+    public void setSingleUseAgents(boolean singleUseAgents) {
+        this.singleUseAgents = singleUseAgents;
+    }
+
     @Override
     public void start(AzureVMComputer azureComputer) {
         //TODO: check when this method is getting called and add code accordingly
@@ -178,7 +237,6 @@ public DescriptorImpl getDescriptor() {
     public static final DescriptorImpl DESCRIPTOR = new DescriptorImpl();
 
     public static class DescriptorImpl extends Descriptor<RetentionStrategy<?>> {
-
         @Override @NonNull
         public String getDisplayName() {
             return "Azure VM Pool Retention Strategy";

src/main/java/com/microsoft/azure/vmagent/builders/AzureVMTemplateFluent.java

@@ -163,9 +163,12 @@ public T addNewIdleRetentionStrategy(String retentionTime) {
         return (T) this;
     }
 
-    public T addNewPoolRetentionStrategy(String retentionTime, String poolSize) {
-        this.retentionStrategy = new AzureVMCloudPoolRetentionStrategy(Integer.parseInt(retentionTime),
+    public T addNewPoolRetentionStrategy(String retentionTime, String poolSize, boolean singleUseAgents) {
+        AzureVMCloudPoolRetentionStrategy retentionStrategy1 =
+          new AzureVMCloudPoolRetentionStrategy(Integer.parseInt(retentionTime),
                 Integer.parseInt(poolSize));
+        retentionStrategy1.setSingleUseAgents(singleUseAgents);
+        this.retentionStrategy = retentionStrategy1;
         return (T) this;
     }
 

src/main/resources/com/microsoft/azure/vmagent/AzureVMCloudPoolRetentionStrategy/config.jelly

@@ -9,4 +9,8 @@
         <f:textbox default="3" />
     </f:entry>
 
+    <f:entry title="${%Single_Use_Agents}" field="singleUseAgents">
+        <f:checkbox/>
+    </f:entry>
+
 </j:jelly>

src/main/resources/com/microsoft/azure/vmagent/AzureVMCloudPoolRetentionStrategy/config.properties

@@ -1,2 +1,3 @@
 RetentionTimeInHours=Retention Time In Hour
 Pool_Size=Pool Size
+Single_Use_Agents=Single Use Agents

src/main/resources/com/microsoft/azure/vmagent/AzureVMCloudPoolRetentionStrategy/help-singleUseAgents.html

@@ -0,0 +1,3 @@
+<div>
+    If enabled, agents are deleted after a single job is run. It is suggested to set retention time to 0 if this is enabled.
+</div>