#!/bin/bash

# Copyright 2014-2023 Google LLC
# test bitmap generation based on ZFS send stream

usage()
{
    cat <<EOF
Usage: $0 [-vhHEVbp] 
	-v - verbose
	-h - help
	-H - test holes
	-E - test embedded data feature
	-n number - number of test iterations, default 1
	-d - debug mode, leave the bitmap and the pool artifacts for further inspection (forces n = 1)
	-V size (bytes) - volume size; default 10M
	-b size (bytes) - blocks size; default 64k
	-g size (bytes) - grain size; default 128k
	-p poolname - name of the pool, default test_pool
	-D "device spec for pool" - quoted list of devices to use, default "/tmp/vd1 /tmp/vd2"
	-Z "path to zfstool" - which tool to execize, default /sbin/zfstool
EOF
}

# Configuration
ntests=1
debug=0
test_holes=0
test_embedded=0
volsize=$((10*1024*1024))
blocksize=$((64*1024))
grainsize=$((64*1024))
verbose=0
poolname="test_pool"
devspec="/tmp/vd1 /tmp/vd2"
zfstool="/sbin/zfstool"

while getopts "hvdHEV:b:g:n:p:D:Z:" opt; do
    case $opt in
	v)
	    verbose=$((verbose+1))
	    ;;
	V)
	    volsize=$OPTARG
	    ;;
	b)
	    blocksize=$OPTARG
	    ;;
	g)
	    grainsize=$OPTARG
	    ;;
	p)
	    poolname=$OPTARG
	    ;;
	H)
	    test_holes=1
	    ;;
	E)
	    test_embedded=1
	    ;;
	D)
	    devspec="$OPTARG"
	    ;;
	Z)
	    zfstool=$OPTARG
	    ;;
	n)
	    ntests=$OPTARG
	    ;;
	d)
	    debug=1
	    ;;
	h)
	    usage
	    exit 0
	    ;;
	\?)
	    echo "Invalid option: -$OPTARG" >&2
	    usage
	    exit -1
	    ;;
    esac
done

if [ $((volsize % blocksize)) -ne 0 ]; then
    echo "Volume size $volsize is not a multiple of blocks of size $blocksize"
    exit -1
fi
if [ $((blocksize % grainsize)) -ne 0 ]; then
    echo "Block size $blocksize is not a multiple of grains of size $grainsize"
    exit -1
fi

if [ $debug -eq 1 ]; then
    n=1
fi

echo "Test configuration:"
echo "The zfstool: $zfstool"
echo "Number of test iterations: $ntests"
echo -n "Debug mode: "
if [ $debug -eq 1 ]; then echo "on"; else echo "off"; fi
echo "Pool name: $poolname"
echo "Device list for the pool: $devspec"
echo "Verbosity level: $verbose"
echo -n "Test holes: "
if [ $test_holes -eq 1 ]; then echo "yes"; else echo "no"; fi
echo -n "Test embedded data pointer: "
if [ $test_embedded -eq 1 ]; then echo "yes"; else echo "no"; fi
echo "Volume size: $((volsize/1024/1024))M"
echo "Block size: $((blocksize/1024))K"
echo "Grain size: $((grainsize/1024))K"

# Globals
volblocks=$((volsize/blocksize))
maxlen=$((volblocks/5))
randsource=`mktemp`
compsource=`mktemp`
zvol_name="$poolname/zvol1"

write_data()
{
    local src=$1
    local dst=$2

    for i in 1 2 3; do
	maxoff=$((volblocks-maxlen))
	off=$((RANDOM % maxoff))
	len=$((RANDOM % maxlen))
	dd if=$src of=$dst bs=$blocksize count=$len seek=$off > /dev/null 2>&1
    done
}

generate_random_source()
{
    touch $randsource
    truncate -s 0 $randsource
    dd if=/dev/urandom of=$randsource bs=$blocksize count=$maxlen > /dev/null 2>&1
}

cleanup_random_source()
{
    rm -f $randsource
}

generate_comp_source()
{
    temp=`mktemp`
    touch $compsource
    truncate -s 0 $compsource
    truncate -s 0 $temp
    for i in `seq 1 $blocksize`; do
	echo -n 'F' >> $temp
    done
    for i in `seq 1 $maxlen`; do
	dd if=$temp of=$compsource bs=$blocksize count=1 seek=$i > /dev/null 2>&1
    done
    rm -f $temp
}

cleanup_comp_source()
{
    rm -f $compsource
}

generate_sources()
{
    generate_random_source
    generate_comp_source
}

cleanup_sources()
{
    cleanup_random_source
    cleanup_comp_source
}

write_test_data()
{
    write_data $randsource /dev/zvol/$zvol_name
    if [ $test_holes -eq 1 ]; then
	write_data /dev/zero /dev/zvol/$zvol_name
    fi
    if [ $test_embedded -eq 1 ]; then
	write_data $compsource /dev/zvol/$zvol_name
    fi
}

snap_and_clone()
{
    local ind=$1
    zfs snapshot $zvol_name@snap$ind
    zfs clone $zvol_name@snap$ind $zvol_name-clone$ind
}

make_zpool()
{
    if [ -d /proc/spl/kstat/zfs/$poolname ]; then
	zpool destroy -f $poolname
    fi
    zpool create -f $poolname $devspec
}

wait_zvol()
{
    udevadm settle
}

# main script
set -e

echo "Prepare data sources"
cleanup_sources
generate_sources
echo "Make zpool $poolname"
make_zpool

# create a volume with block size, sparse and compression properties to test
# the required features

echo "Generating zvol option string"
zvol_opts="-V$volsize -b$blocksize"
if [ $test_holes -eq 1 ]; then
    zvol_opts="$zvol_opts -s"
fi
if [ $test_embedded -eq 1 ] || [ $test_holes -eq 1 ]; then
    zvol_opts="$zvol_opts -o compression=on"
fi
echo "zvol options: $zvol_opts"

echo "Generating zfstool option string"
zfstool_base_opts="-g $grainsize -z"
for i in `seq 1 $verbose`; do
    zfstool_base_opts="$zfstool_base_opts -v"
done
echo "zfstool options: $zfstool_base_opts"

echo "Create zvol $zvol_name"
zfs create $zvol_opts $zvol_name
wait_zvol $zvol_name

bmap_base="/tmp/bitmap"

echo "Writing initial data to $zvol_name"
write_test_data
echo "Creating initial snapshot and clone of $zvol_name"
snap_and_clone 1

for t in `seq 1 $ntests`; do
    echo "Iteration $t"
    # write some random data and some compressible data if needed
    echo "Writing data to $zvol_name"
    write_test_data
    # create snapshot, clone
    echo "Creating snapshot and clone of $zvol_name"
    snap_and_clone $((t+1))
    # test bitmap generation and validate with bitwise comparison
    echo "Testing bitmap generation"    
    bmap=$bmap_base$((t+1))
    holes_bmap=$bmap_base$t.holes
    zfstool_opts="$zfstool_base_opts -c $holes_bmap"
    #zfstool_opts="$zfstool_base_opts"
    echo "$zfstool $zfstool_opts bitmap $zvol_name-clone$t $zvol_name-clone$((t+1)) $bmap"
    $zfstool $zfstool_opts bitmap $zvol_name-clone$t $zvol_name-clone$((t+1)) $bmap
    echo "Testing bitmap correctness"
    echo "$zfstool $zfstool_opts bitmap-verify $zvol_name-clone$t $zvol_name-clone$((t+1))"
    $zfstool $zfstool_opts bitmap-verify $zvol_name-clone$t $zvol_name-clone$((t+1))
done

# cleanup
if [ $debug -eq 0 ]; then
    rm -f $bmap_base*
    echo "Destroying $zvol_name and its snapshots/clones"
    zfs destroy -R $zvol_name
fi

echo "Done with tests"
cleanup_sources