#! /bin/sh

# Run btrfs balance with reasonable values, starting assuming worst case
# Hide the normal 'all is well' output line, but not errors.
# Accept optional directory, so this can be run from the installer.


# btrfs allocated space threshold used to determine when to balance
# default to 80% of the partition size

percent_thresh=80

# btrfs minimum space threshold used to determine when to balance
# default to 2GB

fixed_thresh=2147483648

# For large disks we can rebalance chunks that are less than 50% used.
# however for smaller disks, the disk usage is probably higher, and we
# should rebalance to 100% on the last balance.

final_balance=50


usagemsg='Usage: [-h] [-cdv] [filesystem_directory]\n'

usage()
{
    printf "$usagemsg"
    exit 1;
}

help()
{
    echo "$usagemsg"
    cat << EOF
    -c: Check health, balance if necessary
    -d: Dry-run; display status messages but no changes
    -h: Display this help message
    -v: Verbose; display status messages
EOF
exit 0
}

invalid()
{
  printf "ERROR: Unrecognized argument: $1 \n" >&2
  usage 1
}

# Parse command line options
# A POSIX variable
OPTIND=1         # Reset in case getopts has been used previously in the shell.

# Initialize our own variables:
check_health=0
dryrun=0
verbose=0

while getopts "cdhv" opt; do
    case "$opt" in
    c)  check_health=1  ;;
    d)  dryrun=1  ;;
    h)  help  ;;
    v)  verbose=1  ;;
    *)  invalid $opt  ;;
    esac
done

shift $((OPTIND-1))

[ "${1:-}" = "--" ] && shift

# Set directory if passed in as trailing parameter.
if [ ! -z "${@}" ]; then
    fs="$*"
else
    # default directory to balance
    fs="/"
fi

# Validate directory does not exist, exit with error
if [ ! -d "$fs" ]; then
    printf "ERROR: Directory $fs does not exist.\n" >&2
    exit 1;
fi

# if performing a dryrun, check health and make verbose
if [ $dryrun -ne 0 ]; then
    check_health=1
    verbose=1
fi

# Executes the btrfs balance operation
balance()
{
    if [ $dryrun -eq 0 ]; then
        [ $verbose -ne 0 ] && printf "\nBalancing %s on directory %s" "$1" "$fs"
        result=$(btrfs balance start --full-balance "$1" "$fs" 2>&1)
        if [ `echo $result | grep -q ERROR` ]; then
            printf "\n%s\n" "$result"
        else
            chunks=$( echo $result | grep -oP '(?<=Done,\ had\ to\ )relocate [0-9]+ out of [0-9]+ chunks')
            [ $verbose -ne 0 ] && printf "\n    %s" "$chunks"
        fi
    else
        printf "    'btrfs balance start --full-balance %s %s'\n" "$1" "$fs"
    fi
}

# iterates over multiple fixed values to balance the data and metadata
btrfs_balance()
{
    if [ $dryrun -ne 0 ]; then
        printf "Dryrun option provided; the following commands would have been executed: \n"
    else
        [ $verbose -ne 0 ] && printf "\nInitiating BTRFS balance operation...\n"
    fi
    for pct in 0 5 10 30; do
        balance -musage="$pct"
        balance -dusage="$pct"
    done
    # metadata greater than 30% not useful
    balance -dusage="$final_balance"
    printf "\n\n"
}

# Checks that the amount of space allocated on the device is less than, the larger of the two
# calculated percent_headroom or fixed_headroom.
allocated_space_check_healthy()
{
    [ $verbose -ne 0 ] && printf "\nPerforming device allocated space check...\n"
    # Calculate the values for the max headroom
    percent_headroom=$(( device_size * percent_thresh / 100 ))
    [ $verbose -ne 0 ] && printf "%20s %12d\n" "${percent_thresh}% threshold:" "$percent_headroom"
    fixed_headroom=$(( device_size - fixed_thresh ))
    [ $verbose -ne 0 ] && printf "%20s %12d\n" "2GB threshold:" "$fixed_headroom"
    [ $verbose -ne 0 ] && printf "%32s\n" "=================================="
    max_headroom=$(( percent_headroom > fixed_headroom ? percent_headroom : fixed_headroom ))
    [ $verbose -ne 0 ] && printf "%25s %12d\n" "Max allocated threshold:" "$max_headroom"
    [ $verbose -ne 0 ] && printf "%25s %12d\n" "Current allocated space:" "$data_size"

    # if the device allocated space >= max headroom, we need to balance
    if [ "$device_allocated" -ge "$max_headroom" ]; then
        [ $verbose -ne 0 ] && printf "\nAllocated space exceeds max threshold.\n"
        # return False
        return 1
    else
        [ $verbose -ne 0 ] && printf "\nAllocated space is within max threshold.\n"
        # return True
        return 0
    fi
}

# Checks if the amount of space allocated for the Data information minus the amount
# of space used for the Data information is greater than the chunk size
data_storage_check_efficient()
{
    [ $verbose -ne 0 ] && printf "\nPerforming data storage efficiency check...\n"

    # Fixed chunk size is 1GiB, percent chunk size is 10% of the partition
    fixed_chunk=1073741824
    [ $verbose -ne 0 ] && printf "%20s %12d\n" "Fixed 1GB Chunk:" "$fixed_chunk"
    percent_chunk=$(( device_size * 10 / 100 ))
    [ $verbose -ne 0 ] && printf "%20s %12d\n" "10% device Chunk:" "$percent_chunk"
    [ $verbose -ne 0 ] && printf "%32s\n" "=================================="
    min_chunk_size=$(( percent_chunk < fixed_chunk ? percent_chunk : fixed_chunk ))
    [ $verbose -ne 0 ] && printf "%25s %12d\n" "Min chunk size threshold:" "$min_chunk_size"
    data_free=$(( data_size - data_used ))
    [ $verbose -ne 0 ] && printf "%25s %12d\n" "Current available space:" "$data_free"

    # if the data free space >= min chunk size, a chunk may be freed by running a rebalance.
    if [ "$data_free" -ge "$min_chunk_size" ]; then
        [ $verbose -ne 0 ] && printf "\nOne or more chunks may be freed by running a rebalance.\n"
        # return True
        return 0
    else
        # indicates that no chunks can be freed by running a rebalance
        [ $verbose -ne 0 ] && printf "\nFree space is less than minimum chunk size.\n"
        [ $dryrun -eq 0 ] && printf "Please delete any unnecessary files and execute this command again.\n"
        # return False
        return 1
    fi
}

# gather btrfs usage data and parse into variables. Than call the check functions
# to measure the health of the btrfs filesystem.
btrfs_healthy()
{
    [ $verbose -ne 0 ] && printf "\nCurrent btrfs filesystem state:\n"
    btrfs_fu=$(btrfs filesystem usage -b "$fs")
    # Device values
    device_size=$( echo $btrfs_fu | grep -oP '(?<=Device size:\ )[0-9]+')
    case "$device_size" in
        [0-9]*) [ $verbose -ne 0 ] && printf "%20s %12d\n" "Device Size:" "$device_size";;
        *) printf "ERROR: Failure determining btrfs filesystem size\n" >&2 ; exit 1 ;;
    esac

    device_allocated=$(echo $btrfs_fu | grep -oP '(?<=Device allocated:\ )[0-9]+')
    case "$device_allocated" in
        [0-9]*) [ $verbose -ne 0 ] && printf "%20s %12d\n" "Device Allocated:" "$device_allocated";;
        *) printf "ERROR: Failure determining btrfs filesystem allocated size\n" >&2 ; exit 1 ;;
    esac

    # Data values
    data_single=$(echo $btrfs_fu | grep -oP '(?<=Data,single:\ )Size:[0-9]+, Used:[0-9]+')
    data_size=$(echo $data_single | grep -oP '(?<=Size:)[0-9]+')
    case "$data_size" in
        [0-9]*) [ $verbose -ne 0 ] && printf "%20s %12d\n" "User data size:" "$data_size";;
        *) printf "ERROR: Failure determining btrfs filesystem data size\n" >&2 ; exit 1 ;;
    esac

    data_used=$(echo $data_single | grep -oP '(?<=,\ Used:)[0-9]+')
    case "$data_used" in
        [0-9]*) [ $verbose -ne 0 ] && printf "%20s %12d\n" "User data used:" "$data_size";;
        *) printf "ERROR: Failure determining btrfs filesystem data used\n" >&2 ; exit 1 ;;
    esac

    # if any of the critical variables are empty, exit with error
    if [ -z "$device_size" ] || [ -z "$device_allocated" ] || [ -z "$data_size" ] || [ -z "$data_used" ] ; then
        printf "ERROR: Failed to parse btrfs filesystem usage!\n"
        exit 1;
    fi

    # using the parsed usage parameters check:
    # If balance is necessary by checking the device space allocation
    # Then, if the storage efficiency indicates a chunk can be freed,
    # Execute the balance

    if allocated_space_check_healthy; then
        [ $verbose -ne 0 ] && printf "\nBTRFS filesystem is healthy\n"
        # return True
        return 0
    elif data_storage_check_efficient; then
        [ $verbose -ne 0 ] && printf "\nBTRFS filesystem requires rebalancing.\n"
        # return False
        return 1
    else
        printf "\nWARNING: BTRFS filesystem nearly full; no chunks can be freed with a balance.\nExiting!\n"
        exit 0
    fi
}


# If the check flag is set, execute the health check and determine if balance is necessary.
if [ "$check_health" -ne 0 ]; then
    if ! btrfs_healthy ; then
        # It is safest for the flash disk to first perform the health checks.
        # If this returns a false result, execute the balance commands.
        [ $verbose -ne 0 ] && printf "Filesystem requires rebalancing; starting BTRFS balance.\n"
        # set final balance to maximize chunky freedom
        final_balance=100
        btrfs_balance
    fi
else
    [ $verbose -ne 0 ] && printf "Executing BTRFS balance without health check."
    btrfs_balance
fi

