#!/bin/sh
# Copyright (C) 2000-2010, Parallels, Inc. All rights reserved.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# 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
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#
# vzmigrate is used for container migration to another node


ACT_SCRIPTS_SFX="start stop mount umount premount postumount"
# blowfish is a fast block cipher, much faster then 3des
SSH_OPTIONS="-c blowfish"
SCP_OPTIONS=$SSH_OPTIONS
RSYNC_OPTIONS="-aH --delete --numeric-ids"
DST_PATH=""
VZCTL=vzctl

online=0
verbose=0
remove_area=1
keep_dst=0
debug=0
confdir="/etc/vz/conf"
vzconf="/etc/vz/vz.conf"
tmpdir="/var/tmp"
act_scripts=

# Errors:
MIG_ERR_USAGE=1
MIG_ERR_VPS_IS_STOPPED=2
MIG_ERR_CANT_CONNECT=4
MIG_ERR_COPY=6
MIG_ERR_START_VPS=7
MIG_ERR_STOP_SOURCE=8
MIG_ERR_EXISTS=9
MIG_ERR_NOEXIST=10
MIG_ERR_IP_IN_USE=12
MIG_ERR_QUOTA=13
MIG_ERR_CHECKPOINT=$MIG_ERR_STOP_SOURCE
MIG_ERR_MOUNT_VPS=$MIG_ERR_START_VPS
MIG_ERR_RESTORE_VPS=$MIG_ERR_START_VPS
MIG_ERR_OVZ_NOT_RUNNING=14
MIG_ERR_APPLY_CONFIG=15

# For local vzctl to work, make sure /usr/sbin is in $PATH
if ! echo ":${PATH}:" | fgrep -q ':/usr/sbin:'; then
	PATH="/usr/sbin:$PATH"
fi

usage() {
	cat >&2 <<EOF
This program is used for container migration to another node.
Usage:
vzmigrate [-r yes|no] [--ssh=<options>] [--keep-dst] [--online] [-v]
	[--dst-path=<path>] destination_address <CTID>
Options:
-r, --remove-area yes|no
	Whether to remove container on source host after successful migration.
--ssh=<ssh options>
	Additional options that will be passed to ssh while establishing
	connection to destination host. Please be careful with options
	passed, DO NOT pass destination hostname.
--rsync=<rsync options>
	Additional options that will be passed to rsync.
--keep-dst
	Do not clean synced destination container private area in case of some
	error. It makes sense to use this option on big container migration to
	avoid re-syncing container private area in case some error
	(on container stop for example) occurs during first migration attempt.
--online
	Perform online (zero-downtime) migration: during the migration the
	container freezes for some time and after the migration it
	keeps working as though nothing has happened.
-v
	Verbose mode. Causes vzmigrate to print debugging messages about
	its progress (including some time statistics). Multiple -v options
	increase the verbosity. The maximum is 3.
--dst-path
        Destination node disk path for the migrated container. By default the
        path will be the same.
Examples:
	Online migration of CT #101 to foo.com:
		vzmigrate --online foo.com 101
	Migration of CT #102 to foo.com with downtime:
		vzmigrate foo.com 102
        Migrate CT #101 to /vz4 path:
                vzmigrate --dst-path=/vz4 101 foo.com
Notes:
	This program uses ssh as a transport layer. You need to put ssh
	public key to destination node and be able to connect without
	entering a password.
EOF
	exit $MIG_ERR_USAGE
}

# Logs message
# There are 3 types of messages:
# 0 - error messages (print to stderr)
# 1 - normal messages (print to stdout)
# 2 - debug messages (print to stdout if in verbose mode)
log () {
	if [ $1 -eq 0 ]; then
		shift
		echo "Error: $*" >&2
	elif [ $1 -eq 1 ]; then
		shift
		echo "$*"
	elif [ $verbose -gt 0 ]; then
		shift
		echo "   $@"
	fi
}

# Executes command and returns result of execution
# There are 2 types of execution:
# 1 - normal execution (all output will be printed)
# 2 - debug execution (output will be printed if verbose mode is set,
#     in other case stdout and stderr redirected to /dev/null)
logexec () {
	if [ $1 -eq 1 -o $verbose -gt 0 ]; then
		shift
		$@
	else
		shift
		$@ >/dev/null 2>&1
	fi
}

undo_conf () {
	$SSH "root@$host" "$VZCTL set $VEID --name '' --save > /dev/null"
	$SSH "root@$host" "rm -f $vpsconf"
}

undo_act_scripts () {
	if [ -n "$act_scripts" ] ; then
		$SSH "root@$host" "rm -f $act_scripts"
	fi
	undo_conf
}

undo_private () {
	if [ $keep_dst -eq 0 ]; then
		$SSH "root@$host" "rm -rf $VE_DST_PRIVATE"
	fi
	undo_act_scripts
}

undo_root () {
	$SSH "root@$host" "rm -rf $VE_DST_ROOT"
	undo_private
}

undo_quota_init () {
	[ "${DISK_QUOTA}" = 'no' ] || $SSH "root@$host" "vzquota drop $VEID"
	undo_root
}

undo_quota_on () {
	[ "${DISK_QUOTA}" = 'no' ] || $SSH "root@$host" "vzquota off $VEID"
	undo_quota_init
}

undo_sync () {
	# Root will be destroyed in undo_root
	undo_quota_on
}

undo_suspend () {
	logexec 2 $VZCTL chkpnt $VEID --resume
	undo_sync
}

undo_dump () {
	if [ $debug -eq 0 ]; then
		rm -f "$VE_DUMPFILE"
	fi
	undo_suspend
}

undo_copy_dump () {
	$SSH "root@$host" "rm -f $VE_DUMPFILE"
	undo_suspend
}

undo_stop () {
	if [ "$state" = "running" ]; then
		$VZCTL start $VEID
	elif [ "$mounted" = "mounted" ]; then
		$VZCTL mount $VEID
	fi
	undo_sync
}

undo_source_stage() {
	if [ $online -eq 1 ]; then
		undo_copy_dump
	else
		undo_stop
	fi
}

undo_quota_dump () {
	rm -f "$VE_QUOTADUMP"
	undo_source_stage
}

undo_copy_quota () {
	$SSH "root@$host" "rm -f $VE_QUOTADUMP"
	undo_quota_dump
}

undo_undump () {
	logexec 2 $SSH root@$host $VZCTL restore $VEID --kill
	undo_copy_quota
}

get_status() {
	exist=$3
	mounted=$4
	state=$5
}

get_time () {
	awk -v t2=$2 -v t1=$1 'BEGIN{print t2-t1}'
}

if [ $# -lt 2 ]; then
	usage
fi
while [ ! -z "$1" ]; do
	case "$1" in
	--online)
		online=1
		;;
	-v)
		verbose=$((verbose+1)) # can just be 'let verbose++' in bash
		;;
	-vv)
		verbose=$((verbose+2))
		;;
	-vvv)
		verbose=$((verbose+3))
		;;
	--remove-area|-r)
		shift
		if [ "$1" = "yes" ]; then
			remove_area=1
		elif [ "$1" = "no" ]; then
			remove_area=0
		else
			usage
		fi
		;;
	--keep-dst)
		keep_dst=1
		;;
	--ssh=*)
		SSH_OPTIONS="$SSH_OPTIONS $(echo $1 | cut -c7-)"
		SCP_OPTIONS="`echo $SSH_OPTIONS | sed 's/-p/-P/1'`"
		;;
	--rsync=*)
		RSYNC_OPTIONS="$RSYNC_OPTIONS $(echo $1 | cut -c9-)"
		;;
        --dst-path=*)
                DST_PATH="$(echo $1 | cut -c12-)"
                ;;
	*)
		break
		;;
	esac
	shift
done

if [ $verbose -gt 1 ]; then
	RSYNC_OPTIONS="$RSYNC_OPTIONS -v"
	VZCTL="$VZCTL --verbose"
fi
if [ $verbose -gt 2 ]; then
	SSH_OPTIONS="$SSH_OPTIONS -v"
	SCP_OPTIONS="$SCP_OPTIONS -v"
fi

RSYNC="rsync $RSYNC_OPTIONS"
SSH="ssh $SSH_OPTIONS"
SCP="scp $SCP_OPTIONS"
export RSYNC_RSH="$SSH"

host=$1
shift
VEID=$1
shift

if [ -z "$host" -o -z "$VEID" -o $# -ne 0 ]; then
	usage
fi

# Support CT names as well
if echo $VEID | egrep -qv '^[[:digit:]]+$'; then
	VEID=$(vzlist -o ctid -H $VEID | tr -d ' ')
	if [ -z "$VEID" ]; then
		# Error message is printed by vzlist to stderr
		exit $MIG_ERR_NOEXIST
	fi
fi

vpsconf="$confdir/$VEID.conf"

if [ ! -r "$vzconf" ]; then
	log 0 "Can't read global config file $vzconf"
	exit $MIG_ERR_NOEXIST
fi

get_status $($VZCTL status $VEID)
if [ "$exist" = "deleted" ]; then
	log 0 "CT #$VEID doesn't exist"
	exit $MIG_ERR_NOEXIST
fi

if [ $online -eq 1 ]; then
	log 1 "Starting online migration of CT $VEID to $host"
else
	log 1 "Starting migration of CT $VEID to $host"
fi

# Try to connect to destination
if ! logexec 2 $SSH -o BatchMode=yes root@$host /bin/true; then
	log 0 "Can't connect to destination address using public key"
	log 0 "Please put your public key to destination node"
	exit $MIG_ERR_CANT_CONNECT
fi

# Check if OpenVZ is running
if ! logexec 2 $SSH -o BatchMode=yes root@$host /etc/init.d/vz status ; then
	log 0 "OpenVZ is not running on the target machine"
	log 0 "Can't continue migration"
	exit $MIG_ERR_OVZ_NOT_RUNNING
fi

# Check if CPT modules are loaded for online migration
if [ $online -eq 1 ]; then
	if [ ! -f /proc/cpt ]; then
		log 0 "vzcpt module is not loaded on the source node"
		log 0 "Can't continue online migration"
		exit $MIG_ERR_OVZ_NOT_RUNNING
	fi
	if ! logexec 2 $SSH -o BatchMode=yes root@$host "test -f /proc/rst";
	then
		log 0 "vzrst module is not loaded on the destination node"
		log 0 "Can't continue online migration"
		exit $MIG_ERR_OVZ_NOT_RUNNING
	fi
fi

dst_exist=$($SSH "root@$host" "$VZCTL status $VEID" | awk '{print $3}')
if [ "$dst_exist" = "exist" ]; then
	log 0 "CT #$VEID already exists on destination node"
	exit $MIG_ERR_EXISTS
fi

if [ $online -eq 1 -a "$state" != "running" ]; then
	log 0 "Can't perform online migration of a stopped container"
	exit $MIG_ERR_VPS_IS_STOPPED
fi

log 2 "Loading $vzconf and $vpsconf files"

. "$vzconf"
. "$vpsconf"
VE_DUMPFILE="$tmpdir/dump.$VEID"
VE_QUOTADUMP="$tmpdir/quotadump.$VEID"

if [ -n "$DST_PATH" ]; then
	new_vpsconf="$confdir/$VEID.conf.new"
	VE_DST_ROOT="$DST_PATH/root/$VEID"
	VE_DST_PRIVATE="$DST_PATH/private/$VEID"
	cat $confdir/$VEID.conf | sed -e 's|^\(VE_ROOT=.*\)|VE_ROOT="'$DST_PATH'/root/$VEID"|' -e 's|^\(VE_PRIVATE=.*\)|VE_PRIVATE="'$DST_PATH'/private/$VEID"|' > $new_vpsconf
else
	VE_DST_ROOT="$VE_ROOT"
	VE_DST_PRIVATE="$VE_PRIVATE"
	new_vpsconf=$vpsconf

fi

log 2 "Check IPs on destination node: $IP_ADDRESS"
for IP in $IP_ADDRESS; do
	if [ $($SSH "root@$host" "grep -c \" $IP \" /proc/vz/veip") -gt 0 ];
	then
		log 0 "IP address $IP already in use on destination node"
		exit $MIG_ERR_IP_IN_USE
	fi
done

log 1 "Preparing remote node"

log 2 "Copying config file"
if ! logexec 2 $SCP $new_vpsconf root@$host:$vpsconf ; then
	log 0 "Failed to copy config file"
	exit $MIG_ERR_COPY
fi

logexec 2 $SSH root@$host $VZCTL set $VEID --applyconfig_map name --save
# vzctl return code 20 or 21 in case of unrecognized option
if [ $? != 20 ] && [ $? != 21 ] && [ $? != 0 ]; then
	log 0 "Failed to apply config on destination node"
	undo_conf
	exit $MIG_ERR_APPLY_CONFIG
fi

for sfx in $ACT_SCRIPTS_SFX; do
	file="$confdir/$VEID.$sfx"
	if [ -f "$file" ]; then
		act_scripts="$act_scripts $file"
	fi
done
if [ -n "$act_scripts" ]; then
	log 2 "Copying action scripts"
	if ! logexec 2 $SCP $act_scripts root@$host:$confdir ; then
		log 0 "Failed to copy action scripts"
		undo_conf
		exit $MIG_ERR_COPY
	fi
fi

log 2 "Creating remote container root dir"
if ! $SSH "root@$host" "mkdir -p $VE_DST_ROOT"; then
	log 0 "Failed to make container root directory"
	undo_act_scripts
	exit $MIG_ERR_COPY
fi

log 2 "Creating remote container private dir"
if ! $SSH "root@$host" "mkdir -p $VE_DST_PRIVATE"; then
	log 0 "Failed to make container private area directory"
	undo_private
	exit $MIG_ERR_COPY
fi

if [ "${DISK_QUOTA}" != "no" ]; then
	log 1 "Initializing remote quota"

	log 2 "Quota init"
	if ! $SSH "root@$host" "$VZCTL quotainit $VEID"; then
		log 0 "Failed to initialize quota"
		undo_root
		exit $MIG_ERR_QUOTA
	fi

	log 2 "Turning remote quota on"
	if ! $SSH "root@$host" "$VZCTL quotaon $VEID"; then
		log 0 "Failed to turn quota on"
		undo_quota_init
		exit $MIG_ERR_QUOTA
	fi
else
	log 2 "VZ disk quota disabled -- skipping quota migration"
fi

log 1 "Syncing private"
$RSYNC "$VE_PRIVATE" "root@$host:${VE_DST_PRIVATE%/*}"
# Ignore rsync error 24 "Partial transfer due to vanished source files"
if [ $? != 24 ] && [ $? != 0 ]; then
	log 0 "Failed to sync container private areas"
	undo_quota_on
	exit $MIG_ERR_COPY
fi

if [ $online -eq 1 ]; then
	log 1 "Live migrating container..."

	log 2 "Suspending container"
	time_suspend=$(date +%s.%N)
	if ! logexec 2 $VZCTL chkpnt $VEID --suspend ; then
		log 0 "Failed to suspend container"
		undo_sync
		exit $MIG_ERR_CHECKPOINT
	fi

	log 2 "Dumping container"
	if ! logexec 2 $VZCTL chkpnt $VEID --dump --dumpfile $VE_DUMPFILE ; then
		log 0 "Failed to dump container"
		undo_suspend
		exit $MIG_ERR_CHECKPOINT
	fi

	log 2 "Copying dumpfile"
	time_copy_dump=$(date +%s.%N)
	if ! logexec 2 $SCP $VE_DUMPFILE root@$host:$VE_DUMPFILE ; then
		log 0 "Failed to copy dump"
		undo_dump
		exit $MIG_ERR_COPY
	fi
else
	if [ "$state" = "running" ]; then
		log 1 "Stopping container"
		if ! logexec 2 $VZCTL stop $VEID ; then
			log 0 "Failed to stop container"
			undo_sync
			exit $MIG_ERR_STOP_SOURCE
		fi
	elif [ "$mounted" = "mounted" ]; then
		log 1 "Unmounting container"
		if ! logexec 2 $VZCTL umount $VEID ; then
			log 0 "Failed to umount container"
			undo_sync
			exit $MIG_ERR_STOP_SOURCE
		fi
	fi
fi

if [ "$state" = "running" ]; then
	log 2 "Syncing private (2nd pass)"
	time_rsync2=$(date +%s.%N)
	if ! $RSYNC "$VE_PRIVATE" "root@$host:${VE_DST_PRIVATE%/*}"; then
		log 0 "Failed to sync container private areas"
		undo_source_stage
		exit $MIG_ERR_COPY
	fi
fi

if [ "${DISK_QUOTA}" != "no" ]; then
	log 1 "Syncing 2nd level quota"

	log 2 "Dumping 2nd level quota"
	time_quota=$(date +%s.%N)
	if ! vzdqdump $VEID -U -G -T > "$VE_QUOTADUMP"; then
		log 0 "Failed to dump 2nd level quota"
		undo_quota_dump
		exit $MIG_ERR_QUOTA
	fi

	log 2 "Copying 2nd level quota"
	if ! logexec 2 $SCP $VE_QUOTADUMP root@$host:$VE_QUOTADUMP ; then
		log 0 "Failed to copy 2nd level quota dump"
		undo_quota_dump
		exit $MIG_ERR_COPY
	fi

	log 2 "Load 2nd level quota"
	if ! $SSH "root@$host" "(vzdqload $VEID -U -G -T < $VE_QUOTADUMP &&
			vzquota reload2 $VEID)"; then
		log 0 "Failed to load 2nd level quota"
		undo_copy_quota
		exit $MIG_ERR_QUOTA
	fi
else
	log 2 "VZ disk quota disabled -- skipping quota migration"
fi

if [ $online -eq 1 ]; then
	log 2 "Undumping container"
	time_undump=$(date +%s.%N)
	if ! logexec 2 $SSH root@$host $VZCTL restore $VEID --undump \
			--dumpfile $VE_DUMPFILE --skip_arpdetect ; then
		log 0 "Failed to undump container"
		undo_copy_quota
		exit $MIG_ERR_RESTORE_VPS
	fi

	log 2 "Resuming container"
	if ! logexec 2 $SSH root@$host $VZCTL restore $VEID --resume ; then
		log 0 "Failed to resume container"
		undo_undump
		exit $MIG_ERR_RESTORE_VPS
	fi
	time_finish=$(date +%s.%N)
	log 2 "Times:"
	log 2 "  Suspend + Dump: " $(get_time $time_suspend $time_copy_dump)
	log 2 "  Copy dump file: " $(get_time $time_copy_dump $time_rsync2)
	log 2 "  Second rsync:   " $(get_time $time_rsync2 $time_quota)
	log 2 "  2nd level quota:" $(get_time $time_quota $time_undump)
	log 2 "  Undump + Resume:" $(get_time $time_undump $time_finish)
	log 2 "Total time: " $(get_time $time_suspend $time_finish)

	log 1 "Cleanup"

	log 2 "Killing container"
	logexec 2 $VZCTL chkpnt $VEID --kill
	logexec 2 $VZCTL umount $VEID

	log 2 "Removing dumpfiles"
	rm -f "$VE_DUMPFILE"
	$SSH "root@$host" "rm -f $VE_DUMPFILE"
else
	if [ "$state" = "running" ]; then
		log 1 "Starting container"
		if ! logexec 2 $SSH root@$host $VZCTL start $VEID ; then
			log 0 "Failed to start container"
			undo_copy_quota
			exit $MIG_ERR_START_VPS
		fi
	elif [ "$mounted" = "mounted" ]; then
		log 1 "Mounting container"
		if ! logexec 2 $SSH root@$host $VZCTL mount $VEID ; then
			log 0 "Failed to mount container"
			undo_copy_quota
			exit $MIG_ERR_MOUNT_VPS
		fi
	elif [ "${DISK_QUOTA}" != "no" ]; then
		log 1 "Turning quota off"
		if ! logexec 2 $SSH root@$host vzquota off $VEID ; then
			log 0 "failed to turn quota off"
			undo_copy_quota
			exit $MIG_ERR_QUOTA
		fi
	fi

	log 1 "Cleanup"
fi

if [ $remove_area -eq 1 ]; then
	log 2 "Destroying container"
	logexec 2 $VZCTL destroy $VEID
else
	# Move config as veid.migrated to allow backward migration
	mv -f $vpsconf $vpsconf.migrated
fi
rm -f $new_vpsconf