#!/bin/bash

HOSTNAME=$(hostname)

ORIGIN_IMG="origin.img"
ORIGIN_IMG_SIZE=2048
ORIGIN_LO_DEV="/dev/loop0"

COW_IMG="cow.img"
COW_IMG_SIZE=512
COW_LO_DEV="/dev/loop"

N_SNAPSHOTS=2
ORIGIN_NAME="testorigin"
SNAPSHOT_NAME="testsnap"
SNAPSHOT_TYPE="P"
CHUNK_SIZE=16

N_SMALL_ITERS=8
N_ITERS=$(($N_SMALL_ITERS*$N_SMALL_ITERS))

READ_DELAY=10

NBD_PROXY="blaa"
NBD_PORT="666"
NBD_DEV="/dev/nbd0"
NBD_SERVER="/root/nbd-2.8.4/nbd-server"
NBD_CLIENT="/root/nbd-2.8.4/nbd-client"

RWSPLIT_NAME="origin-rwsplit"

usage () {
    echo "usage: $(basename $0) [--create] [--remove] [--merge] [--stress]" 2>&1
    exit 1
}

attach_lodev () {
    lodev="$1"
    img="$2"

    if losetup $lodev 2>/dev/null ; then
        echo "$lodev is already attached" 2>&1
        exit 1
    fi

    echo "Attaching $img to $lodev"
    losetup $lodev $img
}

create_loopbacks () {
    if [ -f $ORIGIN_IMG -o -f $COW_IMG ]; then
        echo "$ORIGIN_IMG or $COW_IMG already exists" 2>&1
        exit 1
    fi

    echo "Creating $ORIGIN_IMG (${ORIGIN_IMG_SIZE}M)"
    dd if=/dev/zero of=$ORIGIN_IMG bs=1M seek=$ORIGIN_IMG_SIZE count=0 2>/dev/null

    attach_lodev $ORIGIN_LO_DEV $ORIGIN_IMG

    echo "Creating a ext3 filesystem on $ORIGIN_LO_DEV"
    mke2fs -j $ORIGIN_LO_DEV

    s=1
    while [ $s -le $N_SNAPSHOTS ]; do
      echo "Creating $COW_IMG.$s (${COW_IMG_SIZE}M)"
      dd if=/dev/zero of=$COW_IMG.$s bs=1M seek=$COW_IMG_SIZE count=0 2>/dev/null

      attach_lodev $COW_LO_DEV$s $COW_IMG.$s

      s=$(($s + 1))
    done
}

remove_loopbacks () {
    echo "Detaching loopback devices"

    s=1
    while [ $s -le $N_SNAPSHOTS ]; do
      losetup -d $COW_LO_DEV$s

      if [ -f $COW_IMG.$s ]; then
        echo "Removing $COW_IMG.$s"
        rm -f $COW_IMG.$s
      fi

      s=$(($s + 1))
    done

    losetup -d $ORIGIN_LO_DEV

    if [ -f $ORIGIN_IMG ]; then
        echo "Removing $ORIGIN_IMG"
        rm -f $ORIGIN_IMG
    fi
}

create_rwsplit () {
    use_nbd=$1

    n_sectors=$(blockdev --getsize $ORIGIN_LO_DEV)

    if [ "$use_nbd" ]; then
      echo "Exporting $ORIGIN_LO_DEV over NBD"
      $NBD_SERVER $NBD_PORT $ORIGIN_LO_DEV
      
      echo "Sleeping"
      sleep 2
      
      echo "Setting up NBD loop on $NBD_PROXY"
      ssh -f $NBD_PROXY "modprobe nbd && $NBD_CLIENT $HOSTNAME $NBD_PORT $NBD_DEV && $NBD_SERVER $NBD_PORT $NBD_DEV"
      
      echo "Sleeping"
      sleep 5
      
      echo "Hooking up $NBD_DEV to proxy"
      modprobe nbd && $NBD_CLIENT $NBD_PROXY $NBD_PORT $NBD_DEV
      
      echo "Sleeping"
      sleep 5

      split_tbl="0 $n_sectors rwsplit $NBD_DEV 0 $ORIGIN_LO_DEV 0"

      echo "Setting up rwsplit device (/dev/mapper/$RWSPLIT_NAME)"
      echo $split_tbl | dmsetup create $RWSPLIT_NAME
    else
      split_tbl="0 $n_sectors rwsplit $ORIGIN_LO_DEV 0 $ORIGIN_LO_DEV 0 $READ_DELAY 0"

      echo "Setting up rwsplit device (/dev/mapper/$RWSPLIT_NAME)"
      echo $split_tbl | dmsetup create $RWSPLIT_NAME
    fi
}

remove_rwsplit () {
    echo "Remove rwsplit device (/dev/mapper/$RWSPLIT_NAME)"
    dmsetup remove $RWSPLIT_NAME

    if ps -C $(basename $NBD_SERVER) > /dev/null 2>&1; then
      echo "Disconnecting $NBD_DEV"
      $NBD_CLIENT -d $NBD_DEV && rmmod nbd

      echo "Killing NBD loop on $NBD_PROXY"
      ssh $NBD_PROXY "killall nbd-server && $NBD_CLIENT -d $NBD_DEV && rmmod nbd"

      echo "Killing $ORIGIN_LO_DEV NBD export"
      killall nbd-server
    fi
}

create_snapshots () {
    n_sectors=$(blockdev --getsize $ORIGIN_LO_DEV)

    if [ -b /dev/mapper/$RWSPLIT_NAME ]; then
      origin_dev="/dev/mapper/$RWSPLIT_NAME"
    else
      origin_dev="$ORIGIN_LO_DEV"
    fi

    echo "Creating origin device (/dev/mapper/$ORIGIN_NAME)"
    echo "0 $n_sectors snapshot-origin $origin_dev" | dmsetup create $ORIGIN_NAME

    s=1
    while [ $s -le $N_SNAPSHOTS ]; do
      snapshot_tbl="0 $n_sectors snapshot $origin_dev $COW_LO_DEV$s $SNAPSHOT_TYPE $CHUNK_SIZE"

      echo "Creating snapshot of $origin_dev (/dev/mapper/$SNAPSHOT_NAME$s)"
      echo $snapshot_tbl | dmsetup create $SNAPSHOT_NAME$s

      s=$(($s + 1))
    done
}

remove_snapshots () {
    s=1
    while [ $s -le $N_SNAPSHOTS ]; do
      echo "Removing snapshot device (/dev/mapper/$SNAPSHOT_NAME$s)"
      dmsetup remove $SNAPSHOT_NAME$s

      s=$(($s + 1))
    done

    echo "Removing origin device (/dev/mapper/$ORIGIN_NAME)"
    dmsetup remove $ORIGIN_NAME
}

test_snapshot () {
    ss_dev=$1

    n_sectors=$(blockdev --getsize $ORIGIN_LO_DEV)
    n_chunks=$(($n_sectors / $CHUNK_SIZE))

    echo "Pounding $ss_dev with $N_ITERS chunk reads"

    i=0
    while [ $i -lt $N_ITERS ]; do
      dd iflag=direct if=$ss_dev of=/dev/null count=$CHUNK_SIZE bs=512 skip=$((RANDOM % $n_chunks)) 2>/dev/null &
      i=$(($i + 1))
    done

    wait

    echo "Pounding $ss_dev with $N_ITERS chunk writes"

    exceptions=""

    i=0
    while [ $i -lt $N_ITERS ]; do
      chunk=$((RANDOM % $n_chunks))
      dd oflag=direct if=/dev/zero of=$ss_dev count=$CHUNK_SIZE bs=512 seek=$chunk 2>/dev/null &
      exceptions="$chunk:$exceptions"
      i=$(($i + 1))
    done

    wait

    echo "Reading from all the chunks on $ss_dev we've just written to"

    old_IFS=$IFS;IFS=":"
    for chunk in $exceptions; do
      dd iflag=direct if=$ss_dev of=/dev/null count=$CHUNK_SIZE bs=512 skip=$chunk 2>/dev/null &
    done
    IFS=$old_IFS

    wait

    echo "Simultaneously reading and writing to $N_SMALL_ITERS groups of $N_ITERS chunks on $ss_dev"

    i=0
    while [ $i -lt $N_SMALL_ITERS ]; do
      chunk=$((RANDOM % $n_chunks))
      sector=$(($chunk * $CHUNK_SIZE))

      dd iflag=direct if=$ss_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &
      dd oflag=direct if=/dev/zero of=$ss_dev count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 seek=$sector 2>/dev/null &
      dd iflag=direct if=$ss_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &
      wait

      i=$(($i + 1))
    done
}

test_snapshots() {
    s=1
    while [ $s -le $N_SNAPSHOTS ]; do
      echo "Running tests on /dev/mapper/$SNAPSHOT_NAME$s"
      test_snapshot "/dev/mapper/$SNAPSHOT_NAME$s"
      s=$(($s + 1))
    done
}

test_origin() {
    origin_dev="/dev/mapper/$ORIGIN_NAME"

    n_sectors=$(blockdev --getsize $ORIGIN_LO_DEV)
    n_chunks=$(($n_sectors / $CHUNK_SIZE))

    echo "Pounding $origin_dev with $N_ITERS chunk writes"

    i=0
    while [ $i -lt $N_ITERS ]; do
      chunk=$((RANDOM % $n_chunks))
      j=0
      while [ $j -lt $CHUNK_SIZE ]; do
        dd oflag=direct if=/dev/zero of=$origin_dev count=1 bs=512 seek=$(($chunk + j)) 2>/dev/null &
        dd oflag=direct if=/dev/zero of=$origin_dev count=1 bs=512 seek=$(($chunk + $CHUNK_SIZE - j - 1)) 2>/dev/null &
        j=$(($j + 1))
      done
      i=$(($i + 1))
    done

    wait

    echo "Simultaneously reading and writing to $N_SMALL_ITERS groups of $N_ITERS chunks on $origin_dev"

    i=0
    while [ $i -lt $N_SMALL_ITERS ]; do
      chunk=$((RANDOM % $n_chunks))
      sector=$(($chunk * $CHUNK_SIZE))

      dd iflag=direct if=$origin_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &

      s=1
      while [ $s -le $N_SNAPSHOTS ]; do
        ss_dev="/dev/mapper/$SNAPSHOT_NAME$s"
        dd iflag=direct if=$ss_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &
        s=$(($s + 1))
      done

      dd oflag=direct if=/dev/zero of=$origin_dev count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 seek=$sector 2>/dev/null &

      s=1
      while [ $s -le $N_SNAPSHOTS ]; do
        ss_dev="/dev/mapper/$SNAPSHOT_NAME$s"
        dd iflag=direct if=$ss_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &
        s=$(($s + 1))
      done


      dd iflag=direct if=$origin_dev of=/dev/null count=$(($CHUNK_SIZE * $N_ITERS)) bs=512 skip=$sector 2>/dev/null &
      wait

      i=$(($i + 1))
    done
}

# merge_snapshot () {
#     snapshot_tbl=$(dmsetup table $SNAPSHOT_NAME | sed -e 's/\(.* .* snapshot .* .*\) [PN] \(.*\)/\1 M \2/')
#    
#     echo "Removing original snapshot device (/dev/mapper/$SNAPSHOT_NAME)"
#     dmsetup remove $SNAPSHOT_NAME
# 
#     echo "Creating merged snapshot device (/dev/mapper/$SNAPSHOT_NAME)"
#     echo $snapshot_tbl | dmsetup create $SNAPSHOT_NAME
# }
# 
# prep_stress () {
#     n_cow_chunks=$(dmsetup status $SNAPSHOT_NAME | sed -e 's|.*/\(.*\)|\1|')
# 
#     echo "Filling up the COW device with gunk ($n_cow_chunks chunks)"
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks
# }
# 
# stress_snapshot () {
#     n_cow_chunks=$(dmsetup status $SNAPSHOT_NAME | sed -e 's|.*/\(.*\)|\1|')
# 
#     echo "Running stess test on $SNAPSHOT_NAME"
# 
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks &
#     dd if=/dev/zero of=/dev/mapper/$SNAPSHOT_NAME bs=${CHUNK_SIZE}b count=$n_cow_chunks
# }

do_create=""
do_split=""
do_remove=""
do_test=""
do_merge=""
do_stress=""
with_nbd=""

while [ $# -gt 0 ]; do
    case $1 in
        --create)
            do_create="1"
            ;;

        --split)
            do_split="1"
            ;;

        --with-nbd)
            with_nbd="1"
            ;;

        --remove)
            do_remove="1"
            ;;

        --test)
            do_test="1"
            ;;

#         --merge)
#             do_merge="1"
#             ;;
# 
#         --stress)
#             do_stress="1"
#             ;;
# 
        *)
            usage
            ;;
    esac

    shift
done

if [ -z "$do_create" -a -z "$do_remove" -a -z "$do_test" -a -z "$do_merge" -a -z "$do_stress" ]; then
    usage
fi

if [ "$do_create" ]; then
    create_loopbacks

   if [ "$do_split" ]; then
     create_rwsplit $with_nbd
   fi

   create_snapshots
fi

if [ "$do_test" ]; then
    test_origin
    test_snapshots
fi

# if [ "$do_stress" ]; then
#     prep_stress
# fi
# 
# if [ "$do_merge" -o "$do_stress" ]; then
#     merge_snapshot
# fi
# 
# if [ "$do_stress" ]; then
#     stress_snapshot
# fi

if [ "$do_remove" ]; then
    remove_snapshots
    remove_rwsplit
    remove_loopbacks
fi