#!/bin/bash
# SPDX-License-Identifier: GPL-3.0+
# Copyright (C) 2024 Yu Kuai
#
# Tests for blk-throttle

. common/rc
. common/null_blk
. common/cgroup

THROTL_DIR=$(echo "$TEST_NAME" | tr '/' '_')
THROTL_DEV=dev_nullb
declare THROTL_CLEAR_BASE_SUBTREE_CONTROL_IO
declare THROTL_CLEAR_CGROUP2_DIR_CONTROL_IO

group_requires() {
	_have_root
	_have_null_blk
	_have_kernel_option BLK_DEV_THROTTLING
	_have_cgroup2_controller io
	_have_program bc
}

# Create a new null_blk device, and create a new blk-cgroup for test.
_set_up_throtl() {

	if ! _configure_null_blk $THROTL_DEV "$@" power=1; then
		return 1
	fi

	if ! _init_cgroup2; then
		_exit_null_blk
		return 1
	fi

	THROTL_CLEAR_BASE_SUBTREE_CONTROL_IO=
	THROTL_CLEAR_CGROUP2_DIR_CONTROL_IO=
	if ! grep -q io "$(_cgroup2_base_dir)/cgroup.subtree_control"; then
		echo "+io" > "$(_cgroup2_base_dir)/cgroup.subtree_control"
		THROTL_CLEAR_BASE_SUBTREE_CONTROL_IO=true
	fi
	if ! grep -q io "$CGROUP2_DIR/cgroup.subtree_control"; then
		echo "+io" > "$CGROUP2_DIR/cgroup.subtree_control"
		THROTL_CLEAR_CGROUP2_DIR_CONTROL_IO=true
	fi

	mkdir -p "$CGROUP2_DIR/$THROTL_DIR"
	return 0;
}

_clean_up_throtl() {
	rmdir "$CGROUP2_DIR/$THROTL_DIR"
	if [[ $THROTL_CLEAR_CGROUP2_DIR_CONTROL_IO == true ]]; then
		echo "-io" > "$CGROUP2_DIR/cgroup.subtree_control"
	fi
	if [[ $THROTL_CLEAR_BASE_SUBTREE_CONTROL_IO == true ]]; then
		echo "-io" > "$(_cgroup2_base_dir)/cgroup.subtree_control"
	fi

	_exit_cgroup2
	_exit_null_blk
}

_throtl_set_limits() {
	echo "$(cat /sys/block/$THROTL_DEV/dev) $*" > \
		"$CGROUP2_DIR/$THROTL_DIR/io.max"
}

_throtl_remove_limits() {
	echo "$(cat /sys/block/$THROTL_DEV/dev) rbps=max wbps=max riops=max wiops=max" > \
		"$CGROUP2_DIR/$THROTL_DIR/io.max"
}

_throtl_get_max_io_size() {
	cat "/sys/block/$THROTL_DEV/queue/max_sectors_kb"
}

_throtl_issue_fs_io() {
	local path=$1
	local start_time
	local end_time
	local elapsed

	start_time=$(date +%s.%N)

	if [ "$2" == "read" ]; then
		dd if="${path}" of=/dev/null bs="$3" count="$4" iflag=direct status=none
	elif [ "$2" == "write" ]; then
		dd of="${path}" if=/dev/zero bs="$3" count="$4" oflag=direct conv=fdatasync status=none
	fi

	end_time=$(date +%s.%N)
	elapsed=$(echo "$end_time - $start_time" | bc)
	printf "%.0f\n" "$elapsed"
}

_throtl_issue_io() {
	local start_time
	local end_time
	local elapsed

	start_time=$(date +%s.%N)

	if [ "$1" == "read" ]; then
		dd if=/dev/$THROTL_DEV of=/dev/null bs="$2" count="$3" iflag=direct status=none
	elif [ "$1" == "write" ]; then
		dd of=/dev/$THROTL_DEV if=/dev/zero bs="$2" count="$3" oflag=direct status=none
	fi

	end_time=$(date +%s.%N)
	elapsed=$(echo "$end_time - $start_time" | bc)
	printf "%.0f\n" "$elapsed"
}

# Create an asynchronous thread and bind it to the specified blk-cgroup, issue
# IO and then print time elapsed to the second, blk-throttle limits should be
# set before this function.
_throtl_test_io() {

	{
		local rw=$1
		local bs=$2
		local count=$3

		echo "$BASHPID" > "$CGROUP2_DIR/$THROTL_DIR/cgroup.procs"
		_throtl_issue_io "$rw" "$bs" "$count"
	} &

	wait $!
}
