DistForkCommand.java

package hudson.plugins.distfork;

import hudson.Extension;
import hudson.FilePath;
import hudson.FilePath.TarCompression;
import hudson.Launcher;
import hudson.Util;
import hudson.cli.CLICommand;
import hudson.model.Computer;
import hudson.model.Hudson;
import hudson.model.Label;
import hudson.model.Node;
import hudson.model.Queue;
import hudson.model.Queue.Executable;
import hudson.remoting.Callable;
import hudson.remoting.VirtualChannel;
import hudson.remoting.forward.Forwarder;
import hudson.remoting.forward.ForwarderFactory;
import hudson.remoting.forward.PortForwarder;
import hudson.util.StreamTaskListener;
import org.kohsuke.args4j.Argument;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.Option;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.CancellationException;
import java.util.concurrent.Future;

/**
 * @author Kohsuke Kawaguchi
 */
@Extension
public class DistForkCommand extends CLICommand {

    @Option(name="-l",usage="Label for controlling where to execute this command")
    public String label;

    @Option(name="-n",usage="Human readable name that describe this command. Used in Hudson's UI.")
    public String name;

    @Option(name="-d",usage="Estimated duration of this task in milliseconds, or -1 if unknown")
    public long duration = -1;

    @Argument(handler=RestOfArgumentsHandler.class)
    public List<String> commands = new ArrayList<String>();

    @Option(name="-z",metaVar="FILE",
            usage="Zip/tgz file to be extracted into the target remote machine before execution of the command")
    public String zip;

    @Option(name="-Z",metaVar="FILE",
            usage="Bring back the newly added/updated files in the target remote machine after the end of the command " +
                  "by creating a zip/tgz bundle and place this in the local file system by this name.")
    public String returnZip;

    @Option(name="-e",usage="Environment variables to set to the launched process",metaVar="NAME=VAL")
    public Map<String,String> envs = new HashMap<String,String>();

    @Option(name="-f",usage="Local files to be copied to remote locations before the execution of a task",metaVar="REMOTE=LOCAL")
    public Map<String,String> files = new HashMap<String,String>();

    @Option(name="-F",usage="Remote files to be copied back to local locations after the execution of a task",metaVar="LOCAL=REMOTE")
    public Map<String,String> returnFiles = new HashMap<String,String>();

    @Option(name="-L",usage="Local to remote port forwarding",handler=PortForwardingArgumentHandler.class)
    public List<PortSpec> l2rFowrarding = new ArrayList<PortSpec>();

    @Option(name="-R",usage="Remote to local port forwarding",handler=PortForwardingArgumentHandler.class)
    public List<PortSpec> r2lFowrarding = new ArrayList<PortSpec>();

    public String getShortDescription() {
        return "forks a process on a remote machine and connects to its stdin/stdout";
    }

    protected int run() throws Exception {
        if(commands.isEmpty())
            throw new CmdLineException(null, "No commands are specified");

        Hudson h = Hudson.getInstance();

        Label l = null;
        if (label!=null) {
            l = h.getLabel(label);
            if(l.isEmpty()) {
                stderr.println("No such label: "+label);
                return -1;
            }
        }

        // defaults to the command names
        if (name==null) {
            boolean dots=false;
            if(commands.size()>3) {
                name = Util.join(commands.subList(0,3)," ");
                dots=true;
            }

            name = Util.join(commands," ");
            if(name.length()>80) {
                name=name.substring(0,80);
                dots=true;
            }
            
            if(dots)    name+=" ...";
        }

        final int[] exitCode = new int[]{-1};

        DistForkTask t = new DistForkTask(l, name, duration, new Runnable() {
            public void run() {
                StreamTaskListener listener = new StreamTaskListener(stdout);
                try {
                    Computer c = Computer.currentComputer();
                    Node n = c.getNode();
                    FilePath workDir = n.getRootPath().createTempDir("distfork",null);


                    {// copy over files
                        if(zip!=null) {
                            BufferedInputStream in = new BufferedInputStream(new FilePath(channel, zip).read());
                            if(zip.endsWith(".zip"))
                                workDir.unzipFrom(in);
                            else
                                workDir.untarFrom(in, TarCompression.GZIP);
                        }

                        for (Entry<String, String> e : files.entrySet())
                            new FilePath(channel,e.getValue()).copyToWithPermission(workDir.child(e.getKey()));
                    }

                    List<Closeable> cleanUpList = new ArrayList<Closeable>();
                    setUpPortForwarding(l2rFowrarding,channel,c.getChannel(),cleanUpList);
                    setUpPortForwarding(r2lFowrarding,c.getChannel(),channel,cleanUpList);

                    try {
                        long startTime = c.getChannel().call(new GetSystemTime());
                        Launcher launcher = n.createLauncher(listener);
                        exitCode[0] = launcher.launch().cmds(commands)
                                .stdin(stdin).stdout(stdout).stderr(stderr).pwd(workDir).envs(envs).join();

                        if (!returnFiles.isEmpty() || returnZip!=null) {
                            stderr.println("Copying back files");
                            for (Entry<String, String> e : returnFiles.entrySet())
                                workDir.child(e.getValue()).copyToWithPermission(new FilePath(channel,e.getKey()));

                            if (returnZip!=null) {
                                OutputStream os = new BufferedOutputStream(new FilePath(channel,returnZip).write());
                                try {
                                    RootCutOffFilter scanner = new RootCutOffFilter(new TimestampFilter(startTime));
                                    if(returnZip.endsWith(".zip")) {
                                        workDir.zip(os,scanner);
                                    } else {
                                        workDir.tar(TarCompression.GZIP.compress(os),scanner);
                                    }
                                } finally {
                                    os.close();
                                }
                            }
                        }
                    } finally {
                        workDir.deleteRecursive();
                        for (Closeable cl : cleanUpList)
                            cl.close();
                    }
                } catch (InterruptedException e) {
                    listener.error("Aborted");
                    exitCode[0] = -1;
                } catch (Exception e) {
                    e.printStackTrace(listener.error("Failed to execute a process"));
                    exitCode[0] = -1;
                }
            }

            /**
             * Sets up port-forwarding.
             */
            private void setUpPortForwarding(List<PortSpec> fowrarding, VirtualChannel recv, VirtualChannel send, List<Closeable> cleanUpList) throws IOException, InterruptedException {
                for (PortSpec spec : fowrarding) {
                    Forwarder f = ForwarderFactory.create(send, spec.forwardingHost, spec.forwardingPort);
                    cleanUpList.add(PortForwarder.create(recv,spec.receivingPort, f));
                }
            }
        });

        // run and wait for the completion
        Queue q = h.getQueue();
        Future<Executable> f = q.schedule(t, 0).getFuture();
        try {
            f.get();
        } catch (CancellationException e) {
            stderr.println("Task cancelled");
            return -1;
        } catch (InterruptedException e) {
            // if the command itself is aborted, cancel the execution
            f.cancel(true);
            throw e;
        }

        return exitCode[0];
    }

    /**
     * Obtains the system clock.
     */
    private static final class GetSystemTime implements Callable<Long,RuntimeException> {
        public Long call() {
            return System.currentTimeMillis();
        }

        private static final long serialVersionUID = 1L;
    }
}