#!/bin/bash
###############################################################################
# Copyright 2021 IBM Corp.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
#                            #
###############################################################################
# COMPONENT: gpartvol.sh                                                      #
#                                                                             #
# Grow the root partition on multipath disk of the deployed guest.            #
###############################################################################

function check_partition_layout {
    # Check whether the partition count and layout is supported
	partition_count=`lsblk /dev/mapper/$mpathx -l | awk 'NR>1 && $6=="part" {print $0}' | wc -l`
	if [[ $? -ne 0 ]]; then
	    echo "Failed to get partition number with lsblk /dev/mapper/$mpathx cmd."
	    echo "Aborting extending partition. "
	    exit 1
	fi
	if [[ $partition_count -eq 1 ]]; then
	    echo "One partition found on /dev/mapper/$mpathx."
	elif [[ $partition_count -eq 2 ]]; then
	    echo "Two partitions found on /dev/mapper/$mpathx. Will extend the second partition."
	    # check lvm tool installed
	    output=`which pvdisplay`
	    if [[ $? -ne 0 ]]; then
	        echo "Error: lvm tool is not installed while two partitions found."
	        echo "abort extending partition."
	        exit 1
		fi
		# continue to check the second partition is a LVM physical volume
	    output=`pvdisplay /dev/mapper/${part_prefix}2`
	    rc=$?
	    if [[ $rc -ne 0 ]]; then
	        echo "Error: the second partition is not a LVM partition. output: $output."
	        echo "Aborting extending partition"
	        exit 1
	    fi
	else
	    echo "Error: $partition_count partitions found on /dev/mapper/$mpathx."
	    echo "Only one or two partitions is supported on the root volume, aborting extending partition."
	    exit 1
	fi
}

function resize_single_rootfs {
    # handle the filesystem resize of single-nonlvm-root scenario
    echo "Start to resize root file system."
	# get the root partition device
	if [[ ${mpath:(-1)} == 'p' ]]; then
	    # handle the 36005076802880052a000000000001069p1 case
	    mpathx=$mpath
	fi
	# get the root partition type
	root_partition="/dev/mapper/${mpathx}1"
	rootfs_type=`blkid -o export $root_partition | grep TYPE= | awk -F'=' '{print $2}'`
	if [[ $rootfs_type == 'xfs' ]]; then
		# resize with xfs_growfs
		out=`xfs_growfs / 2>&1`
		rc=$?
		if [[ $rc -ne 0 ]]; then
		    echo "Failed to resize root file system, RC: $rc, Output: $out."
		    exit 1
		else
			echo "Root file system resized successfully."
			exit 0
		fi
	else
		# ext file system, resize fs with resize2fs
		out=`resize2fs /dev/mapper/${mpathx}1 2>&1`
		rc=$?
		if [[ $rc -ne 0 ]]; then
		    echo "Failed to resize root file system, RC: $rc, Output: $out."
		    exit 1
		else
		    # In some scenario, the resize2fs does not recognize the new partition size
		    # Sample: "The filesystem is already 2621184 (4k) blocks long.  Nothing to do!"
		    # So add some retry here.
		    if [[ $out =~ "Nothing to do!" ]]; then
		        echo "Doing some retry for resize2fs."
		        sleepTimes=".001 .01 .1 .5 1 2 3 5 8 15 22 34 60"
		        for seconds in $sleepTimes; do
		            sleep $seconds
		            out=`resize2fs /dev/mapper/${mpathx}1 2>&1`
		            rc=$?
		            if [[ $rc -ne 0 ]]; then
		                echo "Failed to resize root file system, RC: $rc, Output: $out."
		                exit 1
		            fi
		            if [[ ! ($out =~ "Nothing to do!") ]]; then
		                break # successful - leave loop
		            fi
		        done
		        if [[ $out =~ "Nothing to do!" ]]; then
		            echo "Failed to resize root file system!"
		            exit 1
		        fi
		    fi
		    echo "Root file system resized successfully."
		    exit 0
		fi
    fi
}

function extend_lvm {
    # Scenario: 
    # two partitions:
    #     part1 is standard partition
    #     part2 is lvm partition (used as lvm physical volume)
    # Handling:
    #     extend the physical volume to use additional space of part2
    output=`pvresize /dev/mapper/${part_prefix}2`
    rc=$?
    if [[ $rc -ne 0 ]]; then
        echo "Error: Failed to resize physical volume ${part_prefix}2, RC: $rc, Output: $output."
        exit 1
    else
        echo "Physical volume resized successfully."
        exit 0
    fi
}

# check multipathd service running
multipathd_status=`systemctl status multipathd | grep "active (running)"`
if [[ $? -ne 0 ]]; then
    echo "multipathd service is inactive, abort extending the root partition."
    exit 1
fi

# In some scenario, the df output the wrong multipath name
echo "Reloading multipath mapping."
multipath -r &> /dev/null

# check parted tool is installed
output=`which parted`
if [[ $? -ne 0 ]]; then
    echo "parted is not installed, please install this tool in the base image."
    echo "abort extending the root partition."
    exit 1
fi

# get root partition
mpath=`lsblk -s -l -o NAME,MOUNTPOINT | egrep ' /$' | awk '{print $1}'`
rc=$?

# sample mpath
# on RHEL7 and RHEL8
# depends on user_friendly_names value in /etc/multipath.conf
# if value is yes:
#    mpatha1
# else value is no (by defalut):
#    36005076802880052a000000000001069 or
#    36005076802880052a000000000001069p1 (sometimes with an additional 'p' before '1')
if [[ $rc -ne 0 || "$mpath" == "" ]]; then
    echo "Unable to find a multipath root partition with lsblk -s."
    exit 1
fi

# remove the last number 1
mpath=${mpath%1*}
part_prefix=$mpath

# handle the 36005076802880052a000000000001069p1 case
if [[ ${mpath:(-1)} == 'p' ]]; then
    mpathx=${mpath%p*}
else
    mpathx=$mpath
fi
echo "root device found : $mpathx."

# Ensure the $mpathx path exists
if [[ ! -r "/dev/mapper/$mpathx" ]]; then
    echo "/dev/mapper/$mpathx does not exist, abort extending partition."
    exit 1
fi

# check partition number
check_partition_layout

# Check partition size and volume size (only for the one single / scenario)
if [[ $partition_count -eq 1 ]]; then
    # Sample lsblk:
    # lsblk /dev/mapper/mpathb -l
    # NAME    MAJ:MIN RM SIZE RO TYPE  MOUNTPOINT
    # mpathb  253:0    0  10G  0 mpath
    # mpathb1 253:1    0  10G  0 part  /
    part_size=`lsblk /dev/mapper/$mpathx -l | awk 'NR>1 && $6=="part" {print $4}'`
    vol_size=`lsblk /dev/mapper/$mpathx -l | awk 'NR>1 && $6=="mpath" {print $4}'`
    echo "Partition size: $part_size, volume size: $vol_size."
    if [[ "$part_size" == "$vol_size" ]]; then
        echo "Partition size equals to the root volume size, no need to extend partition."
        exit 0
    fi
fi

# sample dm_name: dm-0
dm_name=`readlink /dev/mapper/$mpathx | awk -F'/' '{print $2}'`
echo "Find root partition on dm path: $dm_name"
# get the devices name sdx to be used by growpart
sds=`ls /sys/devices/virtual/block/$dm_name/slaves`
part_idx=$partition_count
# grow partition
success=0
for sdN in $sds
do
    out=`parted -s /dev/$sdN resizepart $part_idx 100% 2>&1`
    if [[ $? -eq 0 ]]; then
        success=1
    else
        echo "/dev/$sdN resize failed with parted command, error: $out."
        continue
    fi
done
if [[ $success == 0 ]]; then
    echo "All slaves of $mpathx failed to be resized, abort extending the root partition."
    exit 1
fi
# continue to resize multipath partition
out=`parted -s /dev/mapper/$mpathx resizepart $part_idx 100% 2>&1`
rc=$?
# On RHEL8 the parted for mpathx would return 1 even if the partition is resized successfully.
# so cann't decide whether this is success or not based on the rc and just print all info out.
echo "Partition extended. RC: $rc, Output: $out."

# partprobe - inform the OS of partition table changes
out=`partprobe 2>&1`
rc=$?
if [[ $rc -ne 0 ]]; then
    echo "Failed to partprobe, RC: $rc, Output: $out."
fi

# resize the / for single root partition scenario
# extend the lvm volume group for LVM scenario
if [[ $partition_count -eq 1 ]]; then
    resize_single_rootfs
else
    extend_lvm
fi