JSCH performance test

gistfile1.java

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Date;
import java.util.Properties;
import java.util.Timer;
import java.util.TimerTask;

import org.apache.commons.lang3.time.StopWatch;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.UserInfo;

public class SCPTest {
    
    public static final String USER = "stevebriskin";
    public static final String FROM_DIR = "/Users/stevebriskin/temp/";
    public static final String TO_DIR = "/Users/stevebriskin/";
    public static final String FILENAME = "zeroedfile.dd";
    public static final String KEY_LOC = "/Users/stevebriskin/.ssh/id_rsa";
    
    static Timer timer;
    static SSHUtils.TransferState state;

    public static void main(String[] args) throws Exception{
        
        int[] buffers = new int[]{
                          20 * 1024,
                          100 * 1024,
                          500 * 1024,
                          1 * 1024 * 1024,
                          2 * 1024 * 1024,
                          5 * 1024 * 1024,
                          10 * 1024 * 1024,
                          15 * 1024 * 1024,
                          20 * 1024 * 1024
        };
                         
        
        for( int buffer : buffers){
        timer = new Timer();
        state = new SSHUtils.TransferState();
        
        timer.scheduleAtFixedRate( new TimerTask(){
            public void run() {
                System.out.println(new Date() + "  Bytes sent: " + state.getBytesSent());
            }
            
        }, 10000, 10000);
        
        File file = new File(TO_DIR + FILENAME);
        file.delete();
        
        StopWatch watch = new StopWatch();
        watch.start();
        SSHUtils.execute();
        SSHUtils.scp( new File(FROM_DIR + FILENAME), TO_DIR, state, buffer);
        watch.stop();
        
        //System.out.println(watch);
        System.out.println("Buffer " + buffer + ", Rate: " + state.getBytesSent() / ((watch.getTime() / 1000)) + "bytes/sec" + ", " + watch);
        
        timer.cancel();
        }
        
        System.exit(1);
    }
}

class SSHUtils {
    static Properties _sshConfig = new Properties();
    static{
        _sshConfig.put("StrictHostKeyChecking", "no");
    }
    
    /**
     * @return null if there was no error. Otherwise the error string that was sent.
     */
    public static void execute() throws IOException, JSchException {
        JSch jsch = null;
        Session session = null;
        Channel channel = null;

        InputStream inputStream = null;
        try{
            jsch = new JSch();

            jsch.addIdentity(SCPTest.KEY_LOC);
            
            session = jsch.getSession(SCPTest.USER, "localhost", 22);
            session.setConfig(_sshConfig);
            //machine.authenticate(jsch, session);

            session.connect();
        }finally{
            if (inputStream != null)
                inputStream.close();
            if (channel != null)
                channel.disconnect();
            if (session != null)
                session.disconnect();
        }
    }

    /**
     * @param machine is the machine to SCP a file to
     * @param localFile is the file to SCP
     * @param transferState is an object a caller can pass in to track
     *                      transfer progress. transferState can be null.
     *
     * @return true if the file was successfully transfered. Otherwise return
     *         false
     */
    public static boolean scp(File localFile, String remoteDirectory, TransferState transferState, int blockSize) 
        throws InterruptedException {
        final int BLOCK_SIZE = blockSize;
        final byte[] bytesToCopy = new byte[BLOCK_SIZE];


        ScpTransfer sender = null;
        try{
            sender = new ScpTransfer(localFile);
        }catch (JSchException jschExc) {
            jschExc.printStackTrace();
            return false;
        }

        sender.initFileTransfer(remoteDirectory, localFile.getName());
        if (sender.didTransferFail()) {
            sender._cleanup();
            return false;
        }

        sender.start();

        FileInputStream fileStreamToCopy = null;
        BufferedInputStream inputStream = null;
        try{
            fileStreamToCopy = new FileInputStream(localFile);
            inputStream = new BufferedInputStream(fileStreamToCopy);

            while (sender.didTransferFail() == false) {
                int bytesRead = inputStream.read(bytesToCopy);
                if (bytesRead == -1) {
                    //We're done reading the file, block until the senders finish sending files
                    sender.inputDone();
                    sender.blockTilDoneSending();

                    break;
                }

                if (transferState != null && transferState.getKillTransfer()) {
                    sender.inputDone();
                    sender.interrupt();
                    break;
                }

                if (transferState != null)
                    transferState.addBytesSent(bytesRead);
                sender.addMoreBytes(bytesToCopy, bytesRead);
            }

            fileStreamToCopy.close();
            inputStream.close();
        }catch (IOException ioException) {
            System.out.println("Received an error closing a file stream that was otherwise successfully SCPed?");
            System.out.println("File: " + localFile.getAbsolutePath());
            ioException.printStackTrace();
        }finally{
            sender._cleanup();
        }

        if (sender.didTransferFail())
            return false;

        return true;
    }

    static class ScpTransfer extends Thread {
        final File _localFile;

        volatile byte[] _buffer = null;
        volatile int _length = 0;

        JSch _jsch = null;
        Session _session = null;
        Channel _channel = null;

        OutputStream _scpOutputStream;
        InputStream _scpInputStream;

        volatile boolean _moreBytesToPush = true;
        volatile boolean _stillSendingBytes = true;

        boolean _transferFailed = false;

        ScpTransfer(File localFile) throws JSchException {
            super("scp -p " + 22 + " " + localFile.getAbsolutePath() + " " +
                  SCPTest.USER + "@localhost:" + localFile.getAbsolutePath());
            _localFile = localFile;

            try{
                _jsch = new JSch();
                _jsch.addIdentity(SCPTest.KEY_LOC);
                _session = _jsch.getSession(SCPTest.USER, "localhost", 22);
                

                _session.setConfig(_sshConfig);
                //UserInfo ui= new MyUserInfo();
               // _session.setUserInfo(ui);
                
            }catch (JSchException jschExc) {
                jschExc.printStackTrace();
                _transferFailed = true;
            }
        }

        public void initFileTransfer(String remoteDirectory, String remoteFilename) {
            String ackMsg = null;
            try{
                _session.connect();
                _channel = _session.openChannel("exec");

                //-p means set a specific lastModified timestamp for the file on the remote server
                //I think...
                String command = "scp -p -t " + remoteDirectory + remoteFilename;
                ((ChannelExec)_channel).setCommand(command);

                _scpOutputStream = _channel.getOutputStream();
                _scpInputStream = _channel.getInputStream();

                _channel.connect();
                _scpOutputStream.flush();
                if ((ackMsg = _checkAck(_scpInputStream)) != null) {
                    System.out.println("Bad ack creating SSH connection:");
                    System.out.println("\t" + ackMsg);
                    _transferFailed = true;
                    return;
                }

                command = "T " + (_localFile.lastModified() / 1000) + " 0" +
                    " " + (_localFile.lastModified() / 1000) + " 0\n";
                _scpOutputStream.write(command.getBytes());
                _scpOutputStream.flush();

                if ((ackMsg = _checkAck(_scpInputStream)) != null) {
                    System.out.println("Bad ack after sending the file's last modified timestamp:");
                    System.out.println("\t" + ackMsg);
                    _transferFailed = true;
                    return;
                }

                command = "C0644 " + _localFile.length() + " " + _localFile.getName() + "\n";
                _scpOutputStream.write(command.getBytes());
                _scpOutputStream.flush();

                if ((ackMsg = _checkAck(_scpInputStream)) != null) {
                    System.out.println("Bad ack after sending the file's size and name:");
                    System.out.println("\t" + ackMsg);
                    _transferFailed = true;
                    return;
                }
            }catch (IOException ioException) {
                ioException.printStackTrace();
                _transferFailed = true;
            }catch (JSchException jschException) {
                jschException.printStackTrace();
                _transferFailed = true;
            }
        }

        public boolean didTransferFail() {
            return _transferFailed;
        }

        public boolean isStillSendingBytes() {
            return _stillSendingBytes;
        }

        public void inputDone() {
            _moreBytesToPush = false;
        }

        void _cleanup() {
            try{
                if (_scpInputStream != null)
                    _scpInputStream.close();
                if (_scpOutputStream != null)
                    _scpOutputStream.close();

                if (_channel != null)
                    _channel.disconnect();
                if (_session != null)
                    _session.disconnect();
            }catch (IOException ioExc) {
                ioExc.printStackTrace();
            }
        }

        public void blockTilDoneSending() throws InterruptedException {
            while (_stillSendingBytes && _transferFailed == false)
                Thread.sleep(10);

        }

        public void addMoreBytes(byte[] nextBytes, int length) throws InterruptedException {
            if (length <= 0) {
                throw new IllegalArgumentException();
            }

            if (nextBytes.length < length) {
                throw new IllegalArgumentException();
            }

            while (hasBuffered()) {
                if (_transferFailed)
                    return;

                //Wait for the previous data to finish sending before adding more
                //(can also use wait() here)
                Thread.sleep(10);
            }

            //Notify the sending portion that there is new data to send
            setBufferAndLength(nextBytes, length);
        }

        private synchronized void setBufferAndLength(final byte[] nextBytes, final int length) {
            _buffer = new byte[length];
            System.arraycopy(nextBytes, 0, _buffer, 0, length);
            _length = length;
        }

        private synchronized void clearBuffer() {
            _buffer = null;
            _length = 0;
        }


        private synchronized boolean hasBuffered() {
            return _length > 0;
        }

        private synchronized int getNumBytesInBuffer() {
            return _length;
        }

        public void run() {
            while (_moreBytesToPush || hasBuffered()) {
                if (!hasBuffered()) {
                    //Wait for more data (can probably use wait() here?)
                    try{
                        Thread.sleep(10);
                        continue;
                    }catch (InterruptedException interrupted) {
                        interrupted.printStackTrace();
                        _transferFailed = true;
                        return;
                    }
                }

                try{
                    _scpOutputStream.write(_buffer, 0, getNumBytesInBuffer());
                }catch (IOException ioException) {
                    ioException.printStackTrace();
                    _transferFailed = true;

                    _cleanup();
                    return;
                }

                //Finished writing bytes to the OutputStream, notify we need more bytes
                clearBuffer();
            }

            if (Thread.interrupted()) {
                _transferFailed = true;
                _stillSendingBytes = false;
                _cleanup();
                return;
            }

            try{
                _scpOutputStream.write(0);
                _scpOutputStream.flush();
                if (_checkAck(_scpInputStream) != null) {
                    System.out.println("Bad ack after sending file");
                    _transferFailed = true;
                }

            }catch (IOException ioException) {
                ioException.printStackTrace();
                _transferFailed = true;
                return;
            }

            _cleanup();

            //Set a flag so the owner can see this file was completely transferred
            _stillSendingBytes = false;
        }
    }

    static String _checkAck(InputStream sshInputStream) throws IOException {
        if (Thread.interrupted())
            return "Thread interrupted";

        int byteValue = sshInputStream.read();
        /*
          byteValue may be 0 for success,
          1 for error,
          2 for fatal error,
          -1
        */
        if (byteValue <= 0)
            return null;

        // ???
        if (byteValue > 2)
            return String.valueOf(byteValue);

        //byteValue of 1 or 2; get the error message.
        StringBuilder message = new StringBuilder();
        for (int chr = sshInputStream.read(); chr != '\n'; chr = sshInputStream.read())
            message.append((char)chr);

        return message.toString();
    }

    public static class TransferState {
        volatile long _bytesSent = 0;
        volatile boolean _kill = false;

        private void addBytesSent(long additionalBytesSent) {
            _bytesSent += additionalBytesSent;
        }

        public long getBytesSent() {
            return _bytesSent;
        }
        
        public void setKillTransfer(boolean kill) {
            _kill = kill;
        }
        
        public boolean getKillTransfer() {
            return _kill;
        }
    }
    
}


class MyUserInfo implements UserInfo {

    @Override
    public String getPassphrase() {
        return "";
    }

    @Override
    public String getPassword() {
        return "";
    }

    @Override
    public boolean promptPassphrase(String arg0) {
        return false;
    }

    @Override
    public boolean promptPassword(String arg0) {
        return false;
    }

    @Override
    public boolean promptYesNo(String arg0) {
        return false;
    }

    @Override
    public void showMessage(String arg0) {
        // TODO Auto-generated method stub

    }

}