Configurable node constraint

pom.xml

@@ -1,13 +1,12 @@
 <project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
   <modelVersion>4.0.0</modelVersion>
-
   <parent>
-    <groupId>org.jvnet.hudson.plugins</groupId>
+    <groupId>org.jenkins-ci.plugins</groupId>
     <artifactId>plugin</artifactId>
-    <version>1.377</version>
-    <relativePath>../pom.xml</relativePath>
+    <version>1.457</version>
   </parent>
 
+  <groupId>org.jenkins-ci.plugins</groupId>
   <artifactId>heavy-job</artifactId>
   <version>1.1-SNAPSHOT</version>
   <packaging>hpi</packaging>

src/main/java/hudson/plugins/heavy_job/HeavyJobProperty.java

@@ -24,32 +24,78 @@
 package hudson.plugins.heavy_job;
 
 import hudson.Extension;
+import hudson.model.BuildListener;
+import hudson.model.Environment;
+import hudson.model.EnvironmentList;
+import hudson.model.JobProperty;
+import hudson.model.JobPropertyDescriptor;
 import hudson.model.AbstractBuild;
 import hudson.model.AbstractProject;
 import hudson.model.Executor;
-import hudson.model.JobProperty;
-import hudson.model.JobPropertyDescriptor;
+import hudson.model.Job;
 import hudson.model.Queue.Executable;
 import hudson.model.Queue.Task;
 import hudson.model.queue.AbstractSubTask;
 import hudson.model.queue.SubTask;
-import org.kohsuke.stapler.DataBoundConstructor;
+import hudson.model.queue.WorkUnit;
+import hudson.model.queue.WorkUnitContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.logging.Logger;
+
+import org.kohsuke.stapler.DataBoundConstructor;
 
 /**
  * Keeps track of the number of executors that need to be consumed for this job.
  *
  * @author Kohsuke Kawaguchi
  */
 public class HeavyJobProperty extends JobProperty<AbstractProject<?,?>> {
+
+    private static final Logger LOGGER = Logger.getLogger(HeavyJobProperty.class.getName());
+
     public final int weight;
+    public final boolean sameNode;
 
     @DataBoundConstructor
-    public HeavyJobProperty(int weight) {
+    public HeavyJobProperty(int weight, boolean sameNode) {
         this.weight = weight;
+        this.sameNode = sameNode;
+    }
+
+    @Override
+    public boolean prebuild(final AbstractBuild<?, ?> build, BuildListener listener) {
+        EnvironmentList environments = build.getEnvironments();
+        environments.add(new Environment() {
+
+            @Override
+            public void buildEnvVars(Map<String, String> env) {
+                Executor executor = Executor.currentExecutor();
+                WorkUnitContext context = executor.getCurrentWorkUnit().context;
+                List<WorkUnit> workUnits = context.getWorkUnits();
+                //Collection<? extends SubTask> subTasks = Tasks.getSubTasksOf(context.task);
+                StringBuilder b = new StringBuilder();
+                for (Iterator<WorkUnit> i = workUnits.iterator(); i.hasNext(); ) {
+                    WorkUnit unit = i.next();
+                    try {
+                        b.append(unit.getExecutor().getOwner().getHostName());
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
+                    if (i.hasNext()) {
+                        b.append(",");
+                    }
+                }
+                LOGGER.info(build.getFullDisplayName() + " executors IP addresses: " + b);
+                env.put("EXECUTORS_IPS", b.toString());
+            }
+
+        });
+        return true;
     }
 
     @Override
@@ -64,7 +110,7 @@ public Executable createExecutable() throws IOException {
                 @Override
                 public Object getSameNodeConstraint() {
                     // must occupy the same node as the project itself
-                    return getProject();
+                    return sameNode ? getProject() : null;
                 }
 
                 @Override
@@ -93,6 +139,11 @@ public static class DescriptorImpl extends JobPropertyDescriptor {
         public String getDisplayName() {
             return Messages.HeavyJobProperty_DisplayName();
         }
+
+        @Override
+        public boolean isApplicable(Class<? extends Job> job) {
+            return true;
+        }
     }
 
     public static class ExecutableImpl implements Executable {
@@ -114,5 +165,9 @@ public AbstractBuild<?,?> getBuild() {
         public void run() {
             // nothing. we just waste time
         }
+
+        public long getEstimatedDuration() {
+            return parent.getEstimatedDuration();
+        }
     }
 }

src/main/resources/hudson/plugins/heavy_job/HeavyJobProperty/config.jelly

@@ -26,4 +26,7 @@ THE SOFTWARE.
   <f:entry title="${%Job Weight}" field="weight">
     <f:textbox default="1" />
   </f:entry>
+  <f:entry title="${%Executors must reside on the same node}" field="sameNode">
+     <f:checkbox default="true" />
+  </f:entry>
 </j:jelly>