#!/bin/bash
#set -x 

setterm -cursor off

stty -echoctl # hide ^C

# function called by trap
on_exit() {
    setterm -cursor on
    echo
    exit $1
}

trap 'on_exit' EXIT HUP INT TERM

get_current_backfilling() {
  # Get Current Remapped
  ceph -s -f json | jq '.["osdmap"]["osdmap"]["num_remapped_pgs"]'
}

sp="/-\|"
sc=0
progress() {
  # Display Progress Bar
  step=$((${2:-64} / 64))
  max=$2
  (( max = max - 1))
  min=$1
  printf '%.0s#' $(seq 0 $step $min)  
  printf "\b${sp:$3:1}"
  printf '%.0s_' $(seq $min $step $max)  
  ((sc = sc + 1))
  ((sc==${#sp})) && sc=0
}


help_msg () {
    printf """
usage: $0 [-h] [-s|--step Step] [--b|--backfilling Max_backfilling] POOL NEW_PG_NUM 
 
Description               Gradually increase the number of PG_NUM without overloading the cluster
 
Optional
  -h, --help              show this help message and exit
  -s Step, --step Step  
                          Increasing PG_NUM by 'Step', minimum and default: 64, if possible power of 2
  -M Max_remapped, --max_remapped Max_remapped            
                          Increasing PG_NUM waits until there is less 'Max_remapped' in progress to resume, default: 4

Positional
  POOL                    Pool to apply
  NEW_PG_NUM              New PG_NUM for pool 'POOL', minimum: 64 and Power of 2
"""
}

REBALANCING_STEP=64
MAX_BACKFILLING=4

if [ $# -lt 2 ]
then
    help_msg
    on_exit 1
fi

while [ $# -ge 2 ]; do
    key="$1"
    case $key in
        -h|--help)
            help_msg
            exit 0
            ;;
        -s|--step)
            if [ $2 -lt 64 ]; then
                echo "Step minimum 64 !"
                exit 1
            else
                REBALANCING_STEP=$2
            fi
            shift
            shift
            ;;
        -M|--max_remapped)
            MAX_BACKFILLING=$2
            shift
            shift
            ;;
        *)
            POOL=$1
            if [[ ! $(($2 & ($2-1))) == 0 ]]; then
                echo "NEW_PG_NUM must be a power of 2" 
                exit 1
            else
                NEW_PG_NUM=$2
            fi
            shift
            shift
    esac
done

get_current_pg_num() {
    # Get Current PG_NUM
    pg=$(ceph osd pool get $POOL pg_num|awk '{print $2}')
    echo $pg
}

healthy_wait() {
    # Wait for Backfilling <= MAX_BACKFILLING
  cbf=$(get_current_backfilling)
  while [ $cbf -gt $MAX_BACKFILLING ]; do
    cu=$((REBALANCING_STEP - (cbf - MAX_BACKFILLING)))
    if [ $cu -lt 0 ]
    then
      cu=0
    fi
    echo -ne "\r[$(progress $cu ${REBALANCING_STEP} ${sc})] ==> Current remapped > Max remapped (${cbf}}>${MAX_BACKFILLING})      "
    cbf=$(get_current_backfilling)
    sleep 2
  done
  echo
}

CURRENT_PG_NUM=$(get_current_pg_num)

if [ $CURRENT_PG_NUM -ge $NEW_PG_NUM ]; then
# If CURRENT_PG_NUM >= NEW_PG_NUM
  echo No change
  exit 0
fi

for count in $(seq $((CURRENT_PG_NUM + REBALANCING_STEP)) $REBALANCING_STEP $NEW_PG_NUM); do
# Increase by REBALANCING_STEP
  healthy_wait
  ceph osd pool set $POOL pg_num $count
  ceph osd pool set $POOL pgp_num $count
  sleep 30
done

if [ $CURRENT_PG_NUM -lt $NEW_PG_NUM ]; then
# CURRENT_PG_NUM < NEW_PG_NUM
  ceph osd pool set $POOL pg_num $NEW_PG_NUM
  ceph osd pool set $POOL pgp_num $NEW_PG_NUM
  sleep 30
fi

MAX_BACKFILLING=0
healthy_wait
echo PG_NUM for pool ${POOL} is set to ${NEW_PG_NUM}

on_exit 0
