#!/bin/sh

set -e
DEV=/dev/sde                   # Will be partitionned with four 50GB partitions
MNT=/mnt                       # Mount point
DATA=$MNT/DATA
NR_CHANGES=1000                # Number of file changed between 2 snapshots

echo "This script will DESTROY all data on $DEV"
echo "It will use $MNT mount point and break everything inside"
echo "The results will be logged into «btrfs-bench_balance.log»"
echo
echo "Answer DESTROYER to proceed"
read ANSWER
test "$ANSWER" = "DESTROYER"



###############################################################################
create_and_mount_btrfs_filesystem() {
    sfdisk $DEV <<%
start=        2048, size=   104857600, type=83
start=   104859648, size=   104857600, type=83
start=   209717248, size=   104857600, type=83
start=   314574848, size=   104857600, type=83
%
    mkfs.btrfs -f ${DEV}1
    mount ${DEV}1 $MNT
}
###############################################################################
fill_btrfs_filesystem() {
    btrfs subvolume create $DATA
    cp -a /usr/. $DATA
}
###############################################################################
print_disk_write() {
    iostat -o JSON $DEV |
python3 -c "
import json
import sys
data = json.loads(sys.stdin.read())
print(data['sysstat']['hosts'][0]['statistics'][0]['disk'][0]['kB_wrtn'])
"
}
###############################################################################
make_snapshots() {
    for L in $(seq $1)
    do
        find $DATA -type f -name "*[!+]" | head -$NR_CHANGES |
        while read LL
        do
            cp "$LL" "$LL+" || true
            rm "$LL"
        done
        btrfs subvolume snapshot -r $DATA $MNT/$L
    done
}
###############################################################################
make_snapshots2() {
    for L in $(seq $1)
    do
        find $DATA -type f | head -$NR_CHANGES |
        while read LL
        do
            echo SOME_BYTES >>"$LL" || true
        done
        btrfs subvolume snapshot -r $DATA $MNT/$L
    done
}
###############################################################################
add_one_disk_and_balance() {
    btrfs device add -f ${DEV}2 $MNT
    btrfs balance start -dconvert=raid1 -mconvert=raid1 $MNT
}
###############################################################################
add_third_disk_and_balance() {
    btrfs device add -f ${DEV}3 $MNT
    btrfs balance start --full-balance $MNT
}
###############################################################################
add_fourth_disk_and_balance() {
    btrfs device add -f ${DEV}4 $MNT
    btrfs balance start --full-balance $MNT
}
###############################################################################
remove_fourth_disk_and_balance() {
    btrfs device remove ${DEV}4 $MNT
    btrfs balance start --full-balance $MNT
}
###############################################################################
remove_third_disk_and_balance() {
    btrfs device remove ${DEV}3 $MNT
    btrfs balance start --full-balance $MNT
}
###############################################################################
remove_second_disk() {
    btrfs balance start -dconvert=single -mconvert=dup $MNT
    btrfs device remove ${DEV}2 $MNT
}
###############################################################################
analyse() {
    echo $(date) $2 >&3
    mount ${DEV}1 $MNT
    START=$(date +%s)
    WRTN=$(print_disk_write)
    $2
    umount $MNT
    SECS=$(expr $(date +%s) - $START)
    MB=$(expr '(' $(print_disk_write) - $WRTN ')' / 1024)
    LINE="$LINE<td>$MB</td><td>$SECS</td>"
    echo $(date) "$MB MB written in $SECS seconds by the balancing ($1 snapshots)" >&3
}
###############################################################################
do_test() {
    echo $(date) create_and_mount_btrfs_filesystem >&3
    create_and_mount_btrfs_filesystem
    echo $(date) fill_btrfs_filesystem >&3
    fill_btrfs_filesystem
    echo $(date) $1 $2 with $NR_CHANGES file changes >&3
    $1 $2
    umount $MNT
    LINE="<tr>"
    analyse $1 add_one_disk_and_balance
    analyse $1 add_third_disk_and_balance
    analyse $1 add_fourth_disk_and_balance
    analyse $1 remove_fourth_disk_and_balance
    analyse $1 remove_third_disk_and_balance
    analyse $1 remove_second_disk
    LINE="$LINE</tr>"
    echo "$LINE" >&3
}
###############################################################################

(
    for S in 8 16 32 64 128 # 1 2 4 8 16 32 64 128
    do
        do_test make_snapshots2 $S
    done
) 3>>btrfs-bench_balance.log


(
    for S in 1 2 4 8 16 32 64
    do
        do_test make_snapshots $S
    done
) 3>>btrfs-bench_balance.log