#!/bin/bash
##
# Copyright Red Hat.
# By Sweet Tea Dorminy, Awez Shaikh, Nikhil Kshirsagar
#
# Licensed under the GPL 2. See LICENSE in this repository.
##
set -e
_disableVDO(){
	dmsetup reload $VDO_VOLUME_NAME --table "0 `blockdev --getsz $VDO_DEVICE` error"
	dmsetup resume $VDO_VOLUME_NAME
}

_enableOriginalVDO(){
	dmsetup reload $VDO_VOLUME_NAME --table "${VDO_TABLE}"
	dmsetup resume $VDO_VOLUME_NAME
}

_cleanup(){
	echo "Error detected, cleaning up..."
	umount $MOUNT_POINT || true
	if [[ -n $VDO_DEPENDENT_DEV_ORIGINAL_TABLE ]]; then
		dmsetup reload $VDO_DEPENDENT_DEV --table "${VDO_DEPENDENT_DEV_ORIGINAL_TABLE}"
		dmsetup resume $VDO_DEPENDENT_DEV
	fi
	if [[ -n $VDO_VOLUME_NAME  ]]; then
		DEVICE_NAME=$VDO_VOLUME_NAME
		dmsetup remove $DEVICE_NAME-merge || true
		dmsetup remove $DEVICE_NAME-origin || true
		dmsetup remove $DEVICE_NAME-snap || true
	fi
	losetup -d $LOOPBACK1 || true
	rm $LOOPBACK_DIR/$DEVICE_NAME-tmp_loopback_file || true
	if [[ -n $VDO_TABLE ]]; then
		_disableVDO
	fi
	DEVICE_NAME=$(basename $VDO_BACKING)
	if [[ -n $DEVICE_NAME ]]; then
		dmsetup remove $DEVICE_NAME-merge || true
		dmsetup remove $DEVICE_NAME-origin || true
		dmsetup remove $DEVICE_NAME-snap || true
	fi
	if [[ -n $VDO_TABLE ]]; then
		_enableOriginalVDO
	fi
	losetup -d $LOOPBACK0 || true 
	rm $LOOPBACK_DIR/$DEVICE_NAME-tmp_loopback_file || true
}

_waitForUserToDeleteStuff(){
	MOUNT_POINT=$1
	ANS='n'
	while [[ $ANS != y ]] ;do
		echo " "
		echo "Please remove some files from $MOUNT_POINT, then proceed"
		echo " "
		echo -n "Proceed? [y/n]: "
		read -n 1 ANS
	done
}

_fstrimAndPrompt(){
	MOUNT_POINT=$1
	local KEEPGOING=true

	while $KEEPGOING; do
		fstrim $MOUNT_POINT || true
		FSOUT=$(echo $?)
		USED=$(vdostats $VDO_VOLUME_NAME | awk 'NR==2 {print $5}' |  sed 's/%//')
		echo "Now down to just ${USED}% used"
		if [[ $FSOUT -ne 0 || $USED == 100 ]];then
			_waitForUserToDeleteStuff "${MOUNT_POINT}"
		else
			KEEPGOING=false
		fi
	done
}

_fstrim(){
	DEVICE=$1
	MOUNT_POINT=$2
	local NUMERATOR=$(dmsetup status $DEVICE | awk '{print $4}' | awk -F "/" '{print $1}')
	local DENOMINATOR=$(dmsetup status $DEVICE | awk '{print $4}' | awk -F "/" '{print $2}')

	if [[ $NUMERATOR -lt $DENOMINATOR ]]; then
		echo "Beginning space reclaim process -- running fstrim..."
		_fstrimAndPrompt $MOUNT_POINT
	else
		echo "No room on snapshot for fstrim!"
	fi
}

_unmount(){
	MOUNT_POINT=$1
	local UMNT=true
	while $UMNT; do
		umount $MOUNT_POINT
		UOUT=$(echo $?)
		if [[ $UOUT -ne 0 ]]; then
			echo "Process still has a open file or directory in $MOUNT_POINT"
			sleep 10
		else
			UMNT=false
		fi
	done
	rmdir $MOUNT_POINT
}

_waitForMerge(){
	DEVICE=$1
	local KEEPGOING=true

	while $KEEPGOING; do
		local NUMERATOR=$(dmsetup status $DEVICE | awk '{print $4}' | awk -F "/" '{print $1}')
		local DENOMINATOR=$(dmsetup status $DEVICE | awk '{print $5}')

		if [[ $NUMERATOR -ne $DENOMINATOR ]];then
			printf "Merging, %u more chunks for %s\n" $((NUMERATOR - DENOMINATOR)) $DEVICE
			sleep 1
		else
			KEEPGOING=false
		fi
	done
}

_mergeSnapshot(){
	DEVICE=$1
	DEVICE_NAME=$(basename $DEVICE)
	dmsetup remove $DEVICE_NAME-origin
	dmsetup suspend $DEVICE_NAME-snap
	#dmsetup create $VDO_VOLUME_NAME --table "$(echo $VDO_TABLE | awk "{\$5=\"${VDO_BACKING}\"; print }")"
	MERGE_TABLE=$(dmsetup table $DEVICE_NAME-snap | awk "{\$3=\"snapshot-merge\"; print }")
	dmsetup create $DEVICE_NAME-merge --table "$MERGE_TABLE"

	_waitForMerge $DEVICE_NAME-merge

	dmsetup remove $DEVICE_NAME-merge
	dmsetup remove $DEVICE_NAME-snap
}

_mergeDataSnap(){
	PARENT=$1
	if [[ -n $VDO_DEPENDENT_DEV ]]; then
		dmsetup suspend $VDO_DEPENDENT_DEV
	fi

	_mergeSnapshot $1

	if [[ -n $VDO_DEPENDENT_DEV ]]; then
		dmsetup reload $VDO_DEPENDENT_DEV --table "${VDO_DEPENDENT_DEV_ORIGINAL_TABLE}"
		dmsetup resume $VDO_DEPENDENT_DEV
	fi

	losetup -d $LOOPBACK1
	rm $LOOPBACK_DIR/$(basename $PARENT)-tmp_loopback_file
}

_mergeBackingSnap(){
	_disableVDO
	_mergeSnapshot $VDO_BACKING
	_enableOriginalVDO

	losetup -d $LOOPBACK0
	rm $LOOPBACK_DIR/$(basename $VDO_BACKING)-tmp_loopback_file
}

_mkloop(){
	DEVICE=$1
	LO_DEV_SIZE=${TMPFILESZ:-$(($(blockdev --getsz $DEVICE)*10/100))}
	DEVICE_NAME=$(basename $DEVICE)
	TMPFS=$(df -k $LOOPBACK_DIR | awk 'NR==2 {print $4}')  
	if [[ TMPFS -lt LO_DEV_SIZE ]]; then          
		echo "Not enough free space for Snapshot"
		echo "Specify LOOPBACK_DIR with free space or smaller TMPFILESZ in kb"
		exit 1
	fi
	truncate -s ${LO_DEV_SIZE}M $LOOPBACK_DIR/$DEVICE_NAME-tmp_loopback_file
	LOOPBACK=$(losetup -f $LOOPBACK_DIR/$DEVICE_NAME-tmp_loopback_file --show)
}

_snap(){
	DEVICE=$1
	DEVICE_NAME=$(basename $DEVICE)
	dmsetup create $DEVICE_NAME-origin --table "0 `blockdev --getsz $DEVICE` snapshot-origin $DEVICE"
	_mkloop $DEVICE
	dmsetup create $DEVICE_NAME-snap --table "0 `blockdev --getsz $DEVICE` snapshot $DEVICE $LOOPBACK PO 4096 2 discard_zeroes_cow discard_passdown_origin"
	SNAP="/dev/mapper/${DEVICE_NAME}-snap"
}   

_insertSnapUnderVDO(){
	VDO_BACKING=$(echo $VDO_TABLE | cut -d' ' -f 5)
	_disableVDO
	VDO_BACKING_NAME=$(basename $VDO_BACKING)
	_snap $VDO_BACKING
	LOOPBACK0=$LOOPBACK
	SNAP_UNDER_VDO=$SNAP
	SNAP_VDO_TABLE=$(echo $VDO_TABLE | awk "{ \$5=\"${SNAP_UNDER_VDO}\"; print  }")
	dmsetup reload $VDO_VOLUME_NAME --table "${SNAP_VDO_TABLE}"
	dmsetup resume $VDO_VOLUME_NAME
}

_addSnapAboveVDO(){
	_snap $VDO_DEVICE
	SNAP_OVER_VDO=$SNAP
	LOOPBACK1=$LOOPBACK
}

_tmpMount(){
	DEVICE=$1
	MOUNT_POINT=$(mktemp --tmpdir -d vdo-recover-XXXXXXXX)
	mount $1 $MOUNT_POINT
	echo $MOUNT_POINT
}

_mountVDOSnap(){
	MOUNT=$(_tmpMount $SNAP_OVER_VDO)

	_fstrim $SNAP_OVER_VDO $MOUNT

	echo "Beginning commit of data changes"

	_unmount $MOUNT
}

_repointUpperDevicesOrMountVDO(){
	# Check whether some other device is using VDO. If so, change that
	# device to point at the VDO snap and prompt the user to clean up;
	# else mount the VDO snap, fstrim, &c.
	SNAP_OVER_BASENAME=$(basename $SNAP_OVER_VDO)
	VDO_MAJMIN=$(dmsetup ls | grep \\\b$VDO_VOLUME_NAME\\\s | cut -f 2)
        VDO_MAJMIN_DEPS_VERSION=$(echo $VDO_MAJMIN | sed "s/:/, /g")
	ORIGIN_OVER_BASENAME=$(echo $SNAP_OVER_BASENAME | sed 's/snap$/origin/')
	VDO_DEPENDENT_DEV=$(dmsetup deps | grep "${VDO_MAJMIN_DEPS_VERSION}"\
	       		    | grep -v \\\b$ORIGIN_OVER_BASENAME\\\b\
		    	    | grep -v \\\b$SNAP_OVER_BASENAME\\\b\
			    | cut -d':' -f 1)
	if [[ -n $VDO_DEPENDENT_DEV ]]; then
		echo "Detecting dependent device $VDO_DEPENDENT_DEV on $VDO_VOLUME_NAME -- manual intervention will be required"
		VDO_DEPENDENT_DEV_ORIGINAL_TABLE=$(dmsetup table $VDO_DEPENDENT_DEV)
		DEPENDENT_NEW_TABLE=$(echo $VDO_DEPENDENT_DEV_ORIGINAL_TABLE | sed "s#^${VDO_MAJMIN}\$#${SNAP_OVER_VDO}#g; s#^${VDO_DEVICE}\$#${SNAP_OVER_VDO}#g")
		dmsetup reload $VDO_DEPENDENT_DEV --table "${DEPENDENT_NEW_TABLE}"
		dmsetup resume $VDO_DEPENDENT_DEV
		echo "You may want to remount, and run fstrim, on any filesystem"
	        echo "mounted on ${VDO_DEPENDENT_DEV}."
		_fstrimAndPrompt "the device or filesystem atop ${VDO_DEPENDENT_DEV}"
	else
		_mountVDOSnap
	fi
}

_recoveryProcess(){

	echo "Recovery process started"

	LOOPBACK_DIR=${LOOPBACK_DIR:-$(mktemp -d --tmpdir vdo-loopback-XXX)}

	_insertSnapUnderVDO

	_addSnapAboveVDO

	_repointUpperDevicesOrMountVDO

	echo "Beginning commit of data changes"
	_mergeDataSnap $VDO_VOLUME_NAME

	_mergeBackingSnap

	echo "Recovery process completed, $VDO_VOLUME_NAME is ${USED}% Used"
}

#######################################################################
VDO_DEVICE=$1
VDO_VOLUME_NAME=$(basename $VDO_DEVICE)

if [[ -z $1 ]] || [[ $1 == "--help" ]] || [[ $1 == "-h" ]]; then
	echo "Usage: ./vdo_recover {path to vdo device}"
	exit 1
else

	if [[ $EUID -ne 0 ]]; then
		echo "$0: cannot open $VDO_DEVICE: Permission denied" 1>&2
		exit 1
	else
		for entry in $(dmsetup ls --target vdo)
		do
			if [ ${entry[@]} = $VDO_VOLUME_NAME ]; then

				if grep -qs "$VDO_DEVICE$" /proc/self/mounts ; then
					echo "$VDO_VOLUME_NAME appears mounted."
					grep "$VDO_DEVICE" /proc/self/mounts
					exit 1
				else
					trap _cleanup 0
					VDO_TABLE=$(dmsetup table $VDO_VOLUME_NAME)
					_recoveryProcess
					trap - 0
					exit 0
				fi
			else
				echo "$VDO_DEVICE not present"
			fi
		done
		echo "$VDO_DEVICE not detected -- not running?"
		exit 1
	fi
fi
