#!/bin/bash

# lfs_migrate: a simple tool to copy and check files.
#
# To avoid allocating objects on one or more OSTs, they should be
# deactivated on the MDS via "lctl --device {device_number} deactivate",
# where {device_number} is from the output of "lctl dl" on the MDS.
#
# To guard against corruption, the file is compared after migration
# to verify the copy is correct and the file has not been modified.
# This is not a protection against the file being open by another
# process, but it would catch the worst cases of in-use files, but
# to be 100% safe the administrator needs to ensure this is safe.

RSYNC=${RSYNC:-rsync}
LFS_MIGRATE_RSYNC_MODE=${LFS_MIGRATE_RSYNC_MODE:-false}
ECHO=echo
LFS=${LFS:-lfs}
RSYNC_WITH_HLINKS=false
LFS_MIGRATE_TMP=${TMPDIR:-/tmp}
MIGRATED_SET="$(mktemp ${LFS_MIGRATE_TMP}/lfs_migrate.links.XXXXXX)"
NEWNAME=""
REMOVE_FID='s/^\[[0-9a-fx:]*\] //'

add_to_set() {
	local old_fid="$1"
	local path="$2"

	echo -e "$old_fid $path" >> "$MIGRATED_SET"
}

path_in_set() {
	local path="$1"

	sed -e "$REMOVE_FID" $MIGRATED_SET | grep -q "^$path$"
}

old_fid_in_set() {
	local old_fid="$1"

	grep "^\\$old_fid" "$MIGRATED_SET" | head -n 1 |
		sed -e "$REMOVE_FID"
}

usage() {
    cat -- <<USAGE 1>&2
usage: lfs_migrate [--dry-run] [-h] [--no-rsync|--rsync] [-q] [-R] [-s]
		   [-v] [-y] [-0] [FILE|DIR...]
	--dry-run only print the names of files to be migrated
	-h show this usage message
	--no-rsync do not fall back to rsync mode even if lfs migrate fails
	-q run quietly (don't print filenames or status)
	--rsync force rsync mode instead of using lfs migrate
	-R restripe file using default directory striping
	-s skip file data comparison after migrate
	-v show verbose debug messages
	-y answer 'y' to usage question
	-0 input file names on stdin are separated by a null character

The -c <stripe_count> and -S <stripe_size> options may not be specified at
the same time as the -R option.

The --rsync and --no-rsync options may not be specified at the same time.

If a directory is an argument, all files in the directory are migrated.
If no file/directory is given, the file list is read from standard input.

Any arguments that are not explicitly recognized by the script are passed
through to the 'lfs migrate' utility.

Examples:
      lfs_migrate /mnt/lustre/dir
      lfs_migrate -p newpool /mnt/lustre/dir
      lfs find /test -O test-OST0004 -size +4G | lfs_migrate -y
USAGE
    exit 1
}

cleanup() {
	rm -f "$MIGRATED_SET"
	[ -n "$NEWNAME" ] && rm -f "$NEWNAME"
}

trap cleanup EXIT

OPT_CHECK=true
OPT_DEBUG=false
OPT_NO_RSYNC=false
OPT_DRYRUN=false
OPT_YES=false
OPT_RESTRIPE=false
OPT_NULL=false
OPT_PASSTHROUGH=()
STRIPE_COUNT=""
STRIPE_SIZE=""
POOL=""

# Examine any long options and arguments.  getopts does not support long
# options, so they must be stripped out and classified as either options
# for the script, or passed through to "lfs migrate".
LONG_OPT=false
SHORT_OPT=false
OPTS=()

for f in $(seq 1 $#); do
	arg=${!f}
	if [ "${arg:0:2}" = "--" ]; then
		SHORT_OPT=false
		if [ "$arg" = "--block" ]; then
			BLOCK="$arg"
			OPT_YES=true
		elif [ "$arg" = "--non-block" ]; then
			BLOCK="$arg"
		elif [ "$arg" = "--dry-run" ]; then
			OPT_DRYRUN=true
			OPT_YES=true
		elif [ "$arg" = "--rsync" ]; then
			LFS_MIGRATE_RSYNC_MODE=true
		elif [ "$arg" = "--no-rsync" ]; then
			OPT_NO_RSYNC=true
			OPT_YES=true
		else
			LONG_OPT=true
			OPT_PASSTHROUGH+=("$arg")
		fi
	elif [ "${arg:0:1}" = "-" ]; then
		LONG_OPT=false
		if [ "$arg" == "-b" ]; then
			BLOCK="$arg"
		else
			SHORT_OPT=true
			OPTS+=("$arg")
		fi
	elif $LONG_OPT; then
		LONG_OPT=false
		# This will prevent long options from having file name
		# arguments, but allows long options with no arguments to work.
		if [ -f "$arg" -o -d "$arg" ]; then
			OPTS+=("$arg")
		else
			[ "${OPT_PASSTHROUGH[-1]}" = "--stripe-count" ] &&
				STRIPE_COUNT="$arg"
			[ "${OPT_PASSTHROUGH[-1]}" = "--stripe-size" ] &&
				STRIPE_SIZE="$arg"
			[ "${OPT_PASSTHROUGH[-1]}" = "--pool" ] &&
				POOL="$arg"
			OPT_PASSTHROUGH+=("$arg")
		fi
	elif $SHORT_OPT; then
		[ "${OPTS[-1]}" = "-c" ] &&
			STRIPE_COUNT="$arg"
		[ "${OPTS[-1]}" = "-S" ] &&
			STRIPE_SIZE="$arg"
		[ "${OPTS[-1]}" = "-p" ] &&
			POOL="$arg"
		SHORT_OPT=false
		OPTS+=("$arg")
	else
		OPTS+=("$arg")
	fi
done

# Reset the argument list to include only the short options and file names
set -- "${OPTS[@]}"

while getopts ":hlnqRsvy0" opt $*; do
    case $opt in
	h) usage;;
	l) ;; # maintained for backward compatibility
	n) OPT_DRYRUN=true
	   OPT_YES=true
	   echo "$(basename $0): -n deprecated, use --dry-run instead" 1>&2
	   echo "$(basename $0): to specify non-block, use --non-block instead" 1>&2;;
	q) ECHO=:;;
	R) OPT_RESTRIPE=true;;
	s) OPT_CHECK=false;;
	v) OPT_DEBUG=true; ECHO=echo; OPT_PASSTHROUGH+=("-v");;
	y) OPT_YES=true;;
	0) OPT_NULL=true;;
	*) # Pass through any unrecognized options to 'lfs migrate'
	   OPT_PASSTHROUGH+=("-$OPTARG")
	   if [[ ${!OPTIND:0:1} != "-" && ! -f "${!OPTIND}" &&
		 ! -d "${!OPTIND}" ]]; then
		OPT_PASSTHROUGH+=("${!OPTIND}")
		((OPTIND++))
	   fi;;
    esac
done
shift $((OPTIND - 1))

if $OPT_RESTRIPE && [[ "$STRIPE_COUNT" || "$STRIPE_SIZE" ]]; then
	echo "$(basename $0): Options -c <stripe_count> and -S <stripe_size> "\
	"may not be specified at the same time as the -R option." 1>&2
	exit 1
fi

if $LFS_MIGRATE_RSYNC_MODE && $OPT_NO_RSYNC; then
	echo "$(basename $0): Options --rsync and --no-rsync may not be "\
	"specified at the same time." 1>&2
	exit 1
fi

if ! $OPT_YES; then
	echo ""
	echo "lfs_migrate is currently NOT SAFE for moving in-use files." 1>&2
	echo "Use it only when you are sure migrated files are unused." 1>&2
	echo "" 1>&2
	echo "If emptying an OST that is active on the MDS, new files may" 1>&2
	echo "use it.  To stop allocating any new objects on OSTNNNN run:" 1>&2
	echo "  lctl set_param osp.<fsname>-OSTNNNN*.max_create_count=0'" 1>&2
	echo "on each MDS using the OST(s) being emptied." 1>&2
	echo -n "Continue? (y/n) "
	read CHECK
	[ "$CHECK" != "y" -a "$CHECK" != "yes" ] && exit 1
fi

# if rsync has --xattr support, then try to copy the xattrs.
$RSYNC --help 2>&1 | grep -q xattr && RSYNC_OPTS="$RSYNC_OPTS -X"
$RSYNC --help 2>&1 | grep -q acls && RSYNC_OPTS="$RSYNC_OPTS -A"
# If rsync copies lustre xattrs in the future, then we can skip lfs (bug 22189)
strings $(which $RSYNC) 2>&1 | grep -q lustre && LFS=:

# rsync creates its temporary files with lenient permissions, even if
# permissions on the original files are more strict. Tighten umask here
# to avoid the brief window where unprivileged users might be able to
# access the temporary file.
umask 0077

lfs_migrate() {
	while IFS='' read -d '' OLDNAME; do
		local hlinks=()
		local stripe_size="$STRIPE_SIZE"
		local stripe_count="$STRIPE_COUNT"
		local parent_count=""
		local parent_size=""

		$ECHO -n "$OLDNAME: "

		# avoid duplicate stat if possible
		local nlink_type=($(LANG=C stat -c "%h %F" "$OLDNAME" 	\
				 2> /dev/null))

		# skip non-regular files, since they don't have any objects
		# and there is no point in trying to migrate them.
		if [ "${nlink_type[1]}" != "regular" ]; then
			echo -e "\r\e[K$OLDNAME: not a regular file, skipped"
			continue
		fi

		# working out write perms is hard, let the shell do it
		if [ ! -w "$OLDNAME" ]; then
			echo -e "\r\e[K$OLDNAME: no write permission, skipped"
			continue
		fi

		if $OPT_DRYRUN && ! $OPT_DEBUG; then
			$ECHO "dry run, skipped"
			continue
		fi

		# xattrs use absolute file paths, so ensure provided path is
		# also absolute so that the names can be compared
		local oldname_absolute=$(readlink -f "$OLDNAME")
		if [ -z "$oldname_absolute" ]; then
			echo -e "\r\e[K$OLDNAME: cannot resolve full path, skipped"
			continue
		fi
		OLDNAME=$oldname_absolute

		# In the future, the path2fid and fid2path calls below
		# should be replaced with a single call to
		# "lfs path2links" once that command is available.  The logic
		# for detecting unlisted hard links could then be removed.
		local fid=$($LFS path2fid "$OLDNAME" 2> /dev/null)
		if [ $? -ne 0 ]; then
			echo -n "\r\e[K$OLDNAME: cannot determine FID; skipping; "
			echo "is this a Lustre file system?"
			continue
		fi

		if [[ ${nlink_type[0]} -gt 1 ]] || $RSYNC_WITH_HLINKS; then
			# don't migrate a hard link if it was already migrated
			if path_in_set "$OLDNAME"; then
				$ECHO -e "$OLDNAME: already migrated via another hard link"
				continue
			fi

			# There is limited space available in the xattrs
			# to store all of the hard links for a file, so it's
			# possible that $OLDNAME is part of a link set but is
			# not listed in xattrs and therefore not listed as
			# being migrated.
			local migrated=$(old_fid_in_set "$fid")
			if [ -n "$migrated" ]; then
				$ECHO -e "$OLDNAME: already migrated via another hard link"
				if $LFS_MIGRATE_RSYNC_MODE; then
					# Only the rsync case has to relink.
					# The lfs migrate case preserves the
					# inode so the links are already
					# correct.
					[ "$migrated" != "$OLDNAME" ] &&
						ln -f "$migrated" "$OLDNAME"
				fi
				add_to_set "$fid" "$OLDNAME"
				continue;
			fi
		fi

		if $OPT_RESTRIPE; then
			UNLINK=""
		else
		# if rsync copies Lustre xattrs properly in the future
		# (i.e. before the file data, so that it preserves striping)
		# then we don't need to do this getstripe/mktemp stuff.
			UNLINK="-u"

			[ -z "$stripe_count" ] &&
				stripe_count=$($LFS getstripe -c "$OLDNAME" 2> /dev/null)

			[ -z "$stripe_size" ] &&
				stripe_size=$($LFS getstripe -S "$OLDNAME" 2> /dev/null)

			[ -z "$stripe_count" -o -z "$stripe_size" ] && UNLINK=""
		fi

		if $OPT_DEBUG; then
			if $OPT_RESTRIPE; then
				parent_count=$($LFS getstripe -c \
					       $(dirname "$OLDNAME") 2> \
					       /dev/null)
				parent_size=$($LFS getstripe -S \
					      $(dirname "$OLDNAME") 2> \
					      /dev/null)
			fi

			$ECHO -n "stripe" \
				"count=${stripe_count:-$parent_count}," \
				"size=${stripe_size:-$parent_size}," \
				"pool=${POOL:-not in a pool}: "
		fi

		if $OPT_DRYRUN; then
			$ECHO "dry run, skipped"
			continue
		fi

		if [[ "$stripe_count" && -z "$STRIPE_COUNT" ]]; then
			stripe_count="-c $stripe_count"
		else
			stripe_count=""
		fi
		if [[ "$stripe_size" && -z "$STRIPE_SIZE" ]]; then
			stripe_size="-S $stripe_size"
		else
			stripe_size=""
		fi

		# detect other hard links and store them on a global
		# list so we don't re-migrate them
		local mntpoint=$(df -P "$OLDNAME" |
				awk 'NR==2 { print $NF; exit }')
		if [ -z "$mntpoint" ]; then
			echo -e "\r\e[K$OLDNAME: cannot determine mount point; skipped"
			continue
		fi
		hlinks=$($LFS fid2path "$mntpoint" "$fid" 2> /dev/null)
		if [ $? -ne 0 ]; then
			echo -e "\r\e[K$OLDNAME: cannot determine hard link paths, skipped"
			continue
		fi
		hlinks+=("$OLDNAME")

		# first try to migrate via Lustre tools, then fall back to rsync
		if ! $LFS_MIGRATE_RSYNC_MODE; then
			if $LFS migrate "${OPT_PASSTHROUGH[@]}" ${BLOCK} \
			   ${stripe_count} ${stripe_size} "$OLDNAME" &> \
			   /dev/null; then
				$ECHO "done migrate"
				for link in ${hlinks[*]}; do
					add_to_set "$fid" "$link"
				done
				continue
			elif $OPT_NO_RSYNC; then
				echo -e "\r\e[K$OLDNAME: refusing to fall back to rsync, skipped" 1>&2
				continue
			else
				$ECHO -n "falling back to rsync: "
				LFS_MIGRATE_RSYNC_MODE=true
			fi
		fi

		NEWNAME=$(mktemp $UNLINK "$OLDNAME-lfs_migrate.tmp.XXXXXX")
		if [ $? -ne 0 -o -z "$NEWNAME" ]; then
			echo -e "\r\e[K$OLDNAME: can't make temp file, skipped" 1>&2
			continue
		fi

		[ "$UNLINK" ] && $LFS setstripe ${OPT_PASSTHROUGH}	\
				 ${stripe_count} ${stripe_size}		\
				 "$NEWNAME" &> /dev/null

		# we use --inplace, since we created our own temp file already
		if ! $RSYNC -a --inplace $RSYNC_OPTS "$OLDNAME" "$NEWNAME";then
			echo -e "\r\e[K$OLDNAME: copy error, exiting" 1>&2
			exit 4
		fi

		if $OPT_CHECK && ! cmp -s "$OLDNAME" "$NEWNAME"; then
			echo -e "\r\e[K$NEWNAME: compare failed, exiting" 1>&2
			exit 8
		fi

		if ! mv "$NEWNAME" "$OLDNAME"; then
			echo -e "\r\e[K$OLDNAME: rename error, exiting" 1>&2
			exit 12
		fi

		$ECHO "done migrate via rsync"
		for link in ${hlinks[*]}; do
			if [ "$link" != "$OLDNAME" ]; then
				ln -f "$OLDNAME" "$link"
			fi
			add_to_set "$fid" "$link"
		done

		# If the number of hlinks exceeds the space in the xattrs,
		# when the final path is statted it will have a link count
		# of 1 (all other links will point to the new inode).
		# This flag indicates that even paths with a link count of
		# 1 are potentially part of a link set.
		[ ${#hlinks[*]} -gt 1 ] && RSYNC_WITH_HLINKS=true
	done
}

if [ "$#" -eq 0 ]; then
	if $OPT_NULL; then
		lfs_migrate
	else
		tr '\n' '\0' | lfs_migrate
	fi
else
	while [ "$1" ]; do
		if [ -d "$1" ]; then
			$LFS find "$1" -type f -print0
		else
			echo -en "$1\0"
		fi
		shift
	done | lfs_migrate
fi