From dbc338700232b7de21f11fd2da671477bdb48a0e Mon Sep 17 00:00:00 2001
From: Jesse Glick <jglick@cloudbees.com>
Date: Wed, 5 Jan 2022 17:14:48 -0500
Subject: [PATCH] Resolve longstanding tech debt from #104 & #117 about
 `BodyExecution`. When we do not use pickles, there is no need for
 `transient`. Moving the reference to `ExecutorStepExecution` from
 `PlaceholderTask` avoids a memory leak. This is only possible now that
 `Callback` holds the `ExecutorStepExecution`. And we can now properly handle
 `AsynchronousExecution.interrupt`.

---
 .../support/steps/ExecutorStepExecution.java  | 34 +++++++------------
 1 file changed, 13 insertions(+), 21 deletions(-)

diff --git a/src/main/java/org/jenkinsci/plugins/workflow/support/steps/ExecutorStepExecution.java b/src/main/java/org/jenkinsci/plugins/workflow/support/steps/ExecutorStepExecution.java
index 9c0343d6..c5f3b116 100644
--- a/src/main/java/org/jenkinsci/plugins/workflow/support/steps/ExecutorStepExecution.java
+++ b/src/main/java/org/jenkinsci/plugins/workflow/support/steps/ExecutorStepExecution.java
@@ -38,7 +38,6 @@
 import hudson.slaves.WorkspaceList;
 import java.io.IOException;
 import java.io.Serializable;
-import java.lang.ref.WeakReference;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -93,6 +92,11 @@ public class ExecutorStepExecution extends AbstractStepExecutionImpl {
     private final ExecutorStep step;
     private ExecutorStepDynamicContext state;
 
+    /**
+     * Needed for {@link BodyExecution#cancel} in certain scenarios.
+     */
+    private @CheckForNull BodyExecution body;
+
     ExecutorStepExecution(StepContext context, ExecutorStep step) {
         super(context);
         this.step = step;
@@ -263,7 +267,7 @@ public static final class QueueTaskCancelled extends CauseOfInterruption {
                         } catch (Exception x) {
                             LOGGER.log(Level.WARNING, null, x);
                         }
-                        BodyExecution body = task.body != null ? task.body.get() : null;
+                        BodyExecution body = task.execution.body;
                         if (body == null) {
                             listener.getLogger().println("Agent " + node.getNodeName() + " was deleted, but do not have a node body to cancel");
                             continue;
@@ -315,18 +319,6 @@ public static final class PlaceholderTask implements ContinuedTask, Serializable
          */
         private String cookie;
 
-        /**
-         * Needed for {@link BodyExecution#cancel}.
-         * {@code transient} because we cannot save a {@link BodyExecution} in {@link PlaceholderTask}:
-         * {@code ExecutorPickle} is written to the stream first, which holds a {@link PlaceholderTask},
-         * and the {@link BodyExecution} holds {@link PlaceholderTask.Callback} whose {@link WorkspaceList.Lease}
-         * is not processed by {@code WorkspaceListLeasePickle} since pickles are not recursive.
-         * So we make a best effort and only try to cancel a body within the current session.
-         * TODO try to rewrite this mess
-         * @see RemovedNodeListener
-         */
-        private transient @CheckForNull WeakReference<BodyExecution> body;
-
         /** {@link Authentication#getName} of user of build, if known. */
         private final @CheckForNull String auth;
 
@@ -796,6 +788,7 @@ private static final class Callback extends BodyExecutionCallback.TailCall {
                     finish(cookie);
                 }
                 if (execution != null) {
+                    execution.body = null;
                     boolean _stopping = execution.state.task.stopping;
                     execution.state.task.stopping = true;
                     try {
@@ -897,10 +890,10 @@ private final class PlaceholderExecutable implements ContinuableExecutable, Acce
                         listener.getLogger().println("Running on " + ModelHyperlinkNote.encodeTo(node) + " in " + workspace);
                         ExecutorStepDynamicContext state = new ExecutorStepDynamicContext(PlaceholderTask.this, lease, exec);
                         execution.state = state;
-                        body = new WeakReference<>(context.newBodyInvoker()
+                        execution.body = context.newBodyInvoker()
                                 .withContexts(env, state)
                                 .withCallback(new Callback(cookie, execution))
-                                .start());
+                                .start();
                         LOGGER.fine(() -> "started " + cookie + " in " + runId);
                         context.saveState();
                     } else {
@@ -940,13 +933,12 @@ private final class PlaceholderExecutable implements ContinuableExecutable, Acce
                                 return;
                             }
                             LOGGER.fine(() -> "interrupted " + cookie + " in " + runId);
-                            // TODO save the BodyExecution somehow and call .cancel() here; currently we just interrupt the build as a whole:
                             Timer.get().submit(() -> { // JENKINS-46738
-                                Executor masterExecutor = r.getExecutor();
-                                if (masterExecutor != null) {
-                                    masterExecutor.interrupt();
+                                Executor thisExecutor = /* AsynchronousExecution. */ getExecutor();
+                                BodyExecution body = execution.body;
+                                if (body != null) {
+                                    body.cancel(thisExecutor != null ? thisExecutor.getCausesOfInterruption().toArray(new CauseOfInterruption[0]) : new CauseOfInterruption[0]);
                                 } else { // anomalous state; perhaps build already aborted but this was left behind; let user manually cancel executor slot
-                                    Executor thisExecutor = /* AsynchronousExecution. */getExecutor();
                                     if (thisExecutor != null) {
                                         thisExecutor.recordCauseOfInterruption(r, _listener);
                                     }