#!/usr/bin/env ruby

# Transfer a set of changes made to the origin of a snapshot LV to another
# block device, possibly using SSH to send to a remote system.
#
# Usage: Start with lvmsync --help, or read the README for all the gory
# details.
#
# Copyright (C) 2011-2014 Matt Palmer <matt@hezmatt.org>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 3, as published
# by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# `LICENCE` file for more details.
#
require 'optparse'
require 'lvm'
require 'git-version-bump'
require 'open3'

PROTOCOL_VERSION = "lvmsync PROTO[3]"

include LVM::Helpers

def main()
	# Parse me some options
	options = { :rsh => "ssh" }

	OptionParser.new do |opts|
		opts.banner = "Usage: lvmsync [options]"
		opts.separator ""
		opts.separator "    lvmsync [-v|--verbose] [--snapback <file>] <snapshot device> [--stdout | [<desthost>:]<destdevice>]"
		opts.separator "    lvmsync [-v|--verbose] [--snapback <file>] --apply <changes file> <destdevice>"
		opts.separator "    lvmsync <-V|--version>"
		opts.separator ""

		opts.on("--server", "Run in server mode (deprecated; use '--apply -' instead)") do |v|
			options[:server] = true
		end

		opts.on("-v", "--[no-]verbose",
		        "Run verbosely") { |v| $verbose = v }

		opts.on("-d", "--[no-]debug",
		        "Print debugging information") { |v| $debug = v }

		opts.on("-q", "--[no-]quiet",
		        "Run quietly") { |v| $quiet = v }

		opts.on("-b <file>", "--snapback <file>",
		        "Make a backup snapshot file on the destination") do |v|
			options[:snapback] = v
		end

		opts.on("-a", "--apply <file>",
		        "Apply mode: write the contents of a snapback file to a device") do |v|
			options[:apply] = v
		end

		opts.on("-s", "--stdout", "Write output data to stdout rather than another lvmsync process") do |v|
			options[:stdout] = true
		end

		opts.on("-r", "--data-source", "Read data blocks from a block device other than the snapshot origin") do |v|
			options[:source] = v
		end

		opts.on("-e", "--rsh <command-string>", "Use specified command when invoking SSH") do |v|
			options[:rsh] = v
		end

		opts.on("-V", "--version", "Print version of lvmsync") do |v|
			begin
				puts "lvmsync #{GVB.version}"
				exit 0
			rescue GVB::VersionUnobtainable
				fatal "Unable to determine lvmsync version.\n" +
				      "Install lvmsync as a gem, or run it from within a git checkout"
			end
		end
	end.parse!

	if $quiet and ($verbose or $debug)
		fatal "I can't run quietly *and* verbosely at the same time!"
	end

	if options[:apply]
		if ARGV[0].nil?
			fatal "No destination device specified."
		end
		options[:device] = ARGV[0]
		run_apply(options)
	elsif options[:server]
		info "--server is deprecated; please use '--apply -' instead"
		if (ARGV[0].nil?)
			fatal "No destination block device specified.  WTF?"
		end
		options[:apply] = '-'
		options[:device] = ARGV[0]
		run_apply(options)
	else
		if ARGV[0].nil?
			fatal "No snapshot specified.  Exiting.  Do you need --help?"
		end
		options[:snapdev] = ARGV[0]

		if options[:stdout] and options[:snapback]
			fatal "--snapback cannot be used with --stdout"
		end

		if (options[:stdout].nil? and ARGV[1].nil?)
			fatal "No destination specified."
		end
		if options[:stdout].nil?
			dev, host = ARGV[1].split(':', 2).reverse
			options[:desthost] = host
			options[:destdev] = dev
		end

		run_client(options)
	end
end

def run_apply(opts)
	snapfile = opts[:snapback] ? File.open(opts[:snapback], 'w') : nil
	infile   = opts[:apply] == '-' ? $stdin : File.open(opts[:apply], 'r')
	destdev  = opts[:device]

	process_dumpdata(infile, destdev, snapfile, opts)
ensure
	snapfile.close unless snapfile.nil?
	infile.close unless infile.nil? or infile == $stdin
end

def process_dumpdata(instream, destdev, snapback = nil, opts = {})
	handshake = instream.readline.chomp
	unless handshake == PROTOCOL_VERSION
		fatal "Handshake failed; protocol mismatch? (saw '#{handshake}' expected '#{PROTOCOL_VERSION}'"
	end

	snapback.puts handshake if snapback

	verbose "Writing changed data to #{destdev.inspect}"
	File.open(destdev, 'r+') do |dest|
		while header = instream.read(12)
			offset, chunksize = header.unpack("QN")
			offset = ntohq(offset)

			begin
				debug "Seeking to #{offset}"
				dest.seek offset
			rescue Errno::EINVAL
				# In certain rare circumstances, we want to transfer a block
				# device where the destination is smaller than the source (DRBD
				# volumes is the canonical use case).  So, we ignore attempts to
				# seek past the end of the device.  Yes, this may lose data, but
				# if you didn't notice that your dd shit itself, it's unlikely
				# you're going to notice now.

				info "Write occured past end of device"

				# Skip the chunk of data
				instream.read(chunksize)
				# Go to the next chunk
				next
			end

			if snapback
				snapback.write(header)
				snapback.write dest.read(chunksize)
				# Got to back to where we were before, since the read from dest
				# has advanced the file pointer by `chunksize`
				dest.seek offset
			end
			dest.write instream.read(chunksize)
			debug "Wrote #{chunksize} bytes at #{offset}"
		end
	end
end

def run_client(opts)
	snapshot = opts[:snapdev]
	desthost = opts[:desthost]
	destdev = opts[:destdev]
	outfd = nil

	lv = begin
		LVM::LogicalVolume.new(snapshot)
	rescue RuntimeError => e
		fatal "#{snapshot}: could not find logical volume (#{e.message})"
	end

	unless lv.snapshot?
		fatal "#{snapshot}: Not a snapshot device"
	end

	# Since, in principle, we're not supposed to be reading from snapshot
	# devices directly, the kernel makes no attempt to make the device's read
	# cache stay in sync with the actual state of the device.  As a result,
	# we have to manually drop all caches before the data looks consistent.
	# PERFORMANCE WIN!
	File.open("/proc/sys/vm/drop_caches", 'w') { |fd| fd.print "3" }

	snapback = opts[:snapback] ? "--snapback #{opts[:snapback]}" : ''

	source = opts[:source] || lv.origin.path
	verbose "Data source: #{source}"

	if opts[:stdout]
		dump_changes(lv, source, $stdout, opts)
	else
		verbose = $verbose ? '-v' : ''
		debug   = $debug   ? '-d' : ''

		server_cmd = if desthost
			"#{opts[:rsh]} #{desthost} lvmsync --apply - #{snapback} #{verbose} #{debug} #{destdev}"
		else
			"#{$0} --apply - #{snapback} #{verbose} #{destdev}"
		end

		exit_status = nil
		errors = nil

		Open3.popen3(server_cmd) do |stdin_fd, stdout_fd, stderr_fd, wait_thr|
			fds = [stdout_fd, stderr_fd]

			dump_changes(lv, source, stdin_fd, opts) do
				# Remember that this fires between *every* block sent to the
				# receiver, so don't do anything particularly slow in here!
				until (active_fds = IO.select(fds, [], [], 0)).nil?
					active_fds[0].each do |fd|
						begin
							info "\e[2K\rremote:#{fd.readline}"
						rescue EOFError, Errno::EPIPE
							fd.close
							fds.delete(fd)
						end
					end
				end
			end

			stdin_fd.close

			# Read any residual data that might be left in the stdout/stderr of
			# the remote; we've got to do this with a timeout because of
			# OpenSSH, which, when used in ControlMaster ("multiplexing") mode,
			# holds open stderr, meaning that IO.select will never indicate
			# that stderr is finished.
			until (active_fds = IO.select(fds, [], [], 0.1)).nil?
				active_fds[0].each do |fd|
					begin
						info "\e[2K\rremote:#{fd.readline}"
					rescue EOFError, Errno::EPIPE
						fd.close
						fds.delete(fd)
					end
				end
			end
			exit_status = wait_thr.value if wait_thr
		end

		if (exit_status or $?).exitstatus != 0
			fatal "APPLY FAILED."
		end
	end
end

def dump_changes(snapshot, source, outfd, opts)
	outfd.puts PROTOCOL_VERSION

	start_time = Time.now
	xfer_count = 0
	xfer_size  = 0
	total_size = 0
	change_count = snapshot.changes.length

	File.open(source, 'r') do |origindev|
		snapshot.changes.each do |r|
			xfer_count += 1
			chunk_size = r.last - r.first + 1
			xfer_size  += chunk_size

			debug "Sending chunk #{r.to_s}..."

			origindev.seek(r.first, IO::SEEK_SET)

			begin
				outfd.print [htonq(r.first), chunk_size].pack("QN")
				outfd.print origindev.read(chunk_size)
			rescue Errno::EPIPE
				$stderr.puts "Remote prematurely closed the connection"
				yield if block_given?
				return
			end

			# Progress bar!
			if xfer_count % 100 == 50 and !$quiet
				$stderr.printf "\e[2K\rSending chunk %i of %i, %.2fMB/s",
									xfer_count,
									change_count,
									xfer_size / (Time.now - start_time) / 1048576
				$stderr.flush
			end
			yield if block_given?
		end

		origindev.seek(0, IO::SEEK_END)
		total_size = origindev.tell
	end

	unless $quiet
		$stderr.printf "\rTransferred %i bytes in %.2f seconds\n",
		               xfer_size, Time.now - start_time

		$stderr.printf "You transferred your changes %.2fx faster than a full dd!\n",
		               total_size.to_f / xfer_size
	end
end

# Take a device name in any number of different formats and return a [VG, LV] pair.
# Raises ArgumentError if the name couldn't be parsed.
def parse_snapshot_name(origname)
	case origname
		when %r{^/dev/mapper/(.*[^-])-([^-].*)$} then
			[$1, $2]
		when %r{^/dev/([^/]+)/(.+)$} then
			[$1, $2]
		when %r{^([^/]+)/(.*)$} then
			[$1, $2]
		else
			raise ArgumentError,
			      "Could not determine snapshot name and VG from #{origname.inspect}"
	end
end

def debug(s)
	$stderr.puts s if $debug
end

def verbose(s)
	$stderr.puts s if $verbose or $debug
end

def info(s)
	$stderr.puts s unless $quiet
end

def fatal(s, status=1)
	$stderr.puts "FATAL ERROR: #{s}"
	exit status
end

main
