#!/bin/bash -ex

RUN_TIME=300		# approximate duration of run (seconds)

[ $# -eq 1 ] && RUN_TIME="$1"

IMAGE_NAME="image-$$"
IMAGE_SIZE="1024"	# MB

ID_TIMEOUT="10"		# seconds to wait to get rbd id after mapping
ID_DELAY=".1"		# floating-point seconds to delay before rescan

function get_time() {
	date '+%s'
}

function times_up() {
	local end_time="$1"

	test $(get_time) -ge "${end_time}"
}

function _get_id() {
	[ $# -eq 1 ] || exit 99
	local image_name="$1"
	local id=""

	cd /sys/bus/rbd/devices
	for i in *; do
		if [ "$(cat $i/name)" = "${image_name}" ]; then
			id="$i"
			break
		fi
	done
	cd - >/dev/null

	echo $id
	test -n "${id}"		# return code 0 if id was found
}
function get_id() {
	[ $# -eq 1 ] || exit 99
	local image_name="$1"
	local id=""
	local end_time=$(expr $(get_time) + ${ID_TIMEOUT})

	while ! id=$(_get_id "${image_name}") && ! times_up "${end_time}"; do
		echo "get_id: image not mapped; trying again after delay" >&2
		sleep "${ID_DELAY}"
	done

	echo $id
	test -n "${id}"		# return code 0 if id was found
}

function map_unmap() {
	[ $# -eq 1 ] || exit 99
	local image_name="$1"

	rbd map "${image_name}"
	udevadm settle

	RBD_ID=$(get_id "${image_name}")

	rbd unmap "/dev/rbd${RBD_ID}"
	udevadm settle
}

function setup() {
	[ $# -eq 2 ] || exit 99
	local image_name="$1"
	local image_size="$2"

	[ -d /sys/bus/rbd ] || sudo modprobe rbd

	# allow ubuntu user to map/unmap rbd devices
	sudo chown ubuntu /sys/bus/rbd/add
	sudo chown ubuntu /sys/bus/rbd/remove
	rbd create "${image_name}" --size="${image_size}"
}

function cleanup() {
    	# Have to rely on globals for the trap call
	# rbd unmap "/dev/rbd${RBD_ID}"		|| true
	rbd rm "${IMAGE_NAME}"			|| true
	sudo chown root /sys/bus/rbd/remove	|| true
	sudo chown root /sys/bus/rbd/add	|| true
}
trap cleanup EXIT HUP INT

#### Start

setup "${IMAGE_NAME}" "${IMAGE_SIZE}"

COUNT=0
START_TIME=$(get_time)
END_TIME=$(expr $(get_time) + ${RUN_TIME})
while ! times_up "${END_TIME}"; do
	map_unmap "${IMAGE_NAME}"
	COUNT=$(expr $COUNT + 1)
done
ELAPSED=$(expr "$(get_time)" - "${START_TIME}")

echo "${COUNT} iterations completed in ${ELAPSED} seconds"

exit 0