parallel-ssh-copy.rb

#!/usr/bin/env ruby
# coding: utf-8

require 'shellwords'


def die(msg)
    $stderr.puts(msg)
    exit 1
end

def human_size(s)
    suffixes = ['B', 'kB', 'MB', 'GB', 'TB']

    si = 0
    while suffixes[si + 1] && s >= 1000.0
        si += 1
        s /= 1024.0
    end
    "%.2f #{suffixes[si]}" % s
end

if ARGV[0] && ARGV[0][0] == '-'
    connections = ARGV.shift[1..-1].to_i
end

connections = 24 if !connections || connections <= 0

remote_path = ARGV[0]
local_path  = ARGV[1] || '.'

if !remote_path || ['-h', '--help'].include?(remote_path)
    die('Usage: parallel-ssh-copy ' +
        '[-<connection count>] <remote path> [local path]')
end

rsplit = remote_path.split(':')
if rsplit.size < 2
    die("#{remote_path} is not a valid remote path (colon missing)")
end

remote_server = rsplit.shift
remote_path = rsplit * ':'

if File.directory?(local_path)
    local_path = local_path + '/' + File.basename(remote_path)
end

if File.exists?(local_path)
    die("#{local_path} already exists, will not overwrite")
end

stat = `ssh #{remote_server.shellescape} #{"stat #{remote_path.shellescape}" \
                                           .shellescape}`.split("\n")

size_match = /Size:\s*(\d+)/.match(stat[1])
die("Failed to stat #{remote_path}") unless size_match

size = size_match[1].to_i

chunk_size = size / connections
chunk_size = (chunk_size + 1048575) / 1048576 # to MB

# Sanitize connections count
connections = size / (chunk_size * 1048576)
last_chunk_size = size - connections * chunk_size * 1048576
if last_chunk_size > 0
    connections += 1
end

partial_files = connections.times.map { |ci| "#{local_path}-p#{ci}" }
partial_files_esc = partial_files.map { |pf| pf.shellescape }

([local_path] + partial_files).each do |path|
    if !system("touch #{path.shellescape} && rm -f #{path.shellescape}")
        die("Cannot write to #{path}")
    end
end

puts("Launching #{connections} connections...")

children = partial_files_esc.map.with_index do |pfe, ci|
    sleep 0.1

    fork do
        exec("ssh #{remote_server.shellescape} " +
                 ("dd if=#{remote_path.shellescape} bs=1M " +
                  "count=#{chunk_size} skip=#{ci * chunk_size} " +
                  "status=none").shellescape +
             " > #{pfe}")
    end
end

puts("Loading...")

last_done = 0
first_time = Time.now
last_time = first_time

while !children.empty?
    sleep 1

    done = partial_files.map { |file| File.size(file) }.inject(:+)
    diff = done - last_done
    last_done = done

    time = Time.now
    time_diff = time - last_time
    last_time = time

    speed = diff / time_diff.to_f

    print("#{"%6.2f" % (done * 100.0 / size)} %    " +
          "#{human_size(done)} / #{human_size(size)}   " +
          "#{human_size(speed)}/s\e[K\r")

    children.reject! do |child|
        Process.wait(child, Process::WNOHANG)
    end
end

speed = size / (Time.now - first_time).to_f
puts("Completed   " +
     "#{human_size(size)} / #{human_size(size)}   " +
     "#{human_size(speed)}/s\e[K")
puts("Concatenating...")

if !system("cat #{partial_files_esc * ' '} > #{local_path.shellescape}")
    die('Failed to concatenate chunks')
end
system("rm -f #{partial_files_esc * ' '}")