Backport InboundAgentRule updates from #795 and #802 (#805)

jenkinsci · Jul 15, 2024 · b8f51b0 · b8f51b0
1 parent fc00092
commit b8f51b0
Showing 1 changed file with 23 additions and 6 deletions.
diff --git a/src/main/java/org/jvnet/hudson/test/InboundAgentRule.java b/src/main/java/org/jvnet/hudson/test/InboundAgentRule.java
@@ -36,18 +36,20 @@
 import hudson.slaves.JNLPLauncher;
 import hudson.slaves.RetentionStrategy;
 import hudson.slaves.SlaveComputer;
+import hudson.util.ProcessTree;
 import hudson.util.StreamCopyThread;
 import hudson.util.VersionNumber;
 import java.io.File;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
+import java.util.UUID;
 import java.util.jar.JarFile;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -77,7 +79,8 @@ public final class InboundAgentRule extends ExternalResource {
 
     private static final Logger LOGGER = Logger.getLogger(InboundAgentRule.class.getName());
 
-    private final ConcurrentMap<String, Process> procs = new ConcurrentHashMap<>();
+    private final String id = UUID.randomUUID().toString();
+    private final Map<String, Process> procs = Collections.synchronizedMap(new HashMap<>());
 
     /**
      * The options used to (re)start an inbound agent.
@@ -348,6 +351,8 @@ public void start(AgentArguments agentArguments, Options options) throws Excepti
         cmd.addAll(agentArguments.commandLineArgs);
         ProcessBuilder pb = new ProcessBuilder(cmd);
         pb.redirectErrorStream(true);
+        pb.environment().put("INBOUND_AGENT_RULE_ID", id);
+        pb.environment().put("INBOUND_AGENT_RULE_NAME", options.getName());
         LOGGER.info(() -> "Running: " + pb.command());
         Process proc = pb.start();
         procs.put(options.getName(), proc);
@@ -389,10 +394,14 @@ public void stop(@NonNull RealJenkinsRule rjr, @NonNull String name) throws Thro
      * You need only call this to simulate an agent crash, followed by {@link #start}.
      */
     public void stop(@NonNull String name) throws InterruptedException {
-        Process proc = procs.remove(name);
+        Process proc = procs.put(name, null);
         if (proc != null) {
+            LOGGER.info(() -> "Killing " + name + " agent JVM (but not subprocesses)");
             proc.destroyForcibly();
             proc.waitFor();
+        } else {
+            // (normal when called from #start)
+            LOGGER.fine(() -> "No " + name + " agent JVM found to kill");
         }
     }
 
@@ -406,14 +415,22 @@ public boolean isAlive(String name) {
     }
 
     @Override protected void after() {
-        for (String name : procs.keySet()) {
+        for (var entry : procs.entrySet()) {
+            String name = entry.getKey();
+            Process proc = entry.getValue();
             try {
-                stop(name);
+                LOGGER.info(() -> "Cleaning up " + name + " agent JVM and/or any subprocesses");
+                ProcessTree.get().killAll(proc, Map.of("INBOUND_AGENT_RULE_ID", id, "INBOUND_AGENT_RULE_NAME", name));
+                if (proc != null) {
+                    proc.destroyForcibly();
+                    proc.waitFor();
+                }
             } catch (InterruptedException e) {
                 e.printStackTrace();
                 Thread.currentThread().interrupt();
             }
         }
+        procs.clear();
     }
 
     public static class AgentArguments implements Serializable {