#!/bin/bash

set -o errexit
set -o pipefail

MOUNT=${MOUNT:-/mnt/lustre}
DIR=${DIR:-$MOUNT/$HOSTNAME}
SCRIPT=$(realpath "$0")
LUSTRE=${LUSTRE:-$(dirname "$SCRIPT")/..}
. $LUSTRE/tests/test-framework.sh

CHROOT="no"
PREFIX="on"
LIBS71=()
DBENCHPID=""

usage() {
	cat <<-EOF
	Usage: $SCRIPT [-C] [-D DIR] [dbench options]
	-C use chroot instead of change directory
	-D DIR use DIR as work directory
	EOF

	exit 1
}

clean_dbench() {
	local rc=$?

	trap '' ERR HUP INT TERM EXIT

	if (( rc > 128 )); then
		local signal=$((rc - 128))
		echo "dbench killed by signal $signal"
		rc=0
	fi

	echo "stopping dbench $PREFIX $DIR at $(date) with return code $rc"

	if [[ -n "$DBENCHPID" ]]; then
		if ps -h -p "$DBENCHPID" && kill "$DBENCHPID"; then
			echo "killed dbench main pid $DBENCHPID"
		fi
	fi

	if [[ -d "$DIR" ]]; then
		echo "clean dbench files $PREFIX $DIR"
		pushd "$DIR"
		rm -frv clients client.txt dbench "${LIBS71[@]}"
		popd
	fi

	if (( rc == 0 )); then
		echo "dbench successfully finished"
	else
		echo "dbench exit with error $rc"
	fi

	exit $rc
}

trap clean_dbench ERR HUP INT TERM EXIT

if [[ -n "$DBENCH_LIB" ]]; then
	PATH=${DBENCH_LIB}:${PATH}
fi

while getopts "CD:" OPT; do
	case "$OPT" in
		D) test -d "$OPTARG" && DIR="$OPTARG" ;;
		C) CHROOT="yes" ;;
		\?) usage ;;
	esac
done

shift $((OPTIND - 1))

echo "looking for dbench program"
if ! type -P dbench; then
	if [[ -z "$MISSING_DBENCH_OK" ]]; then
		error "dbench is not installed"
	fi

	skip_env "dbench is not installed"

	exit 0
fi

if [[ ! -d "$DIR" ]]; then
	echo "creating output directory $DIR"
	mkdir -pv "$DIR"
fi

if ! cd "$DIR"; then
	error "failed to change directory to $DIR"
fi

if [[ -n "$DBENCH_SRC" && -s "$DBENCH_SRC" ]]; then
	echo "use dbench client file $DBENCH_SRC"
	cp -v "$DBENCH_SRC" client.txt
else
	CLIENT_DIRS=(/usr/share/dbench /usr/local/share /usr/lib/dbench)
	CLIENT_FILES=(client.txt client_plain.txt dbench_client)

	if [[ -n "$DBENCH_LIB" ]]; then
		CLIENT_DIRS=("$DBENCH_LIB" "${CLIENT_DIRS[@]}")
	fi

	for CLIENT_DIR in "${CLIENT_DIRS[@]}"; do
		for CLIENT_FILE in "${CLIENT_FILES[@]}"; do
			CLIENT_PATH="$CLIENT_DIR/$CLIENT_FILE"
			if [[ -s "$CLIENT_PATH" ]]; then
				echo "found dbench client file $CLIENT_PATH"
				cp -v "$CLIENT_PATH" client.txt
				break 2
			fi
		done
	done
fi

if [[ ! -s client.txt ]]; then
	skip_env "no dbench client file found" \
		"DBENCH_LIB=$DBENCH_LIB DBENCH_SRC=$DBENCH_SRC"
	exit 0
fi

if [[ "$CHROOT" == "yes" ]]; then
	echo "copying required dbench files to chroot directory $DIR"
	cp -v "$(type -P dbench)" dbench
	LIBS71=($(ldd dbench | awk '
		{
			if (NF == 2 && $1 ~ /^[/]/)
				lib = $1
			else if (NF == 4 && $3 ~ /^[/]/)
				lib = $3
			else
				next
			print(substr(lib, 2))
		}
	'))
	if (( ${#LIBS71[@]} == 0 )); then
		error "failed to get required dbench libraries"
	fi
	if ! tar -C / -chf - "${LIBS71[@]}" | tar -xvf -; then
		error "failed to copy libs ${LIBS71[*]} to $DIR"
	fi
	RUN="chroot $DIR"
	PREFIX="in"
	PATH=.:/:$PATH
fi

export PATH
echo "running 'dbench $*' $PREFIX $DIR at $(date)"
$RUN dbench -c client.txt "$@" &
DBENCHPID=$!
echo "waiting for dbench pid $DBENCHPID"
wait $DBENCHPID