#! /bin/bash
# SPDX-License-Identifier: GPL-2.0
# Copyright (c) 2023 Meta Platforms, Inc.  All Rights Reserved.
#
# FS QA Test 301
#
# Test common btrfs simple quotas scenarios involving sharing extents and
# removing them in various orders.
#
. ./common/preamble
_begin_fstest auto quick qgroup clone subvol prealloc snapshot remount

. ./common/reflink

_supported_fs btrfs
_require_scratch_reflink
_require_cp_reflink
_require_btrfs_command inspect-internal dump-tree
_require_xfs_io_command "falloc"
_require_scratch_enable_simple_quota
_require_no_compress

subv=$SCRATCH_MNT/subv
nested=$SCRATCH_MNT/subv/nested
snap=$SCRATCH_MNT/snap
nr_fill=512
fill_sz=$((8 * 1024))
total_fill=$(($nr_fill * $fill_sz))
nodesize=$($BTRFS_UTIL_PROG inspect-internal dump-super $SCRATCH_DEV | \
					grep nodesize | $AWK_PROG '{print $2}')
ext_sz=$((128 * 1024 * 1024))
limit_nr=8
limit=$(($ext_sz * $limit_nr))

prep_fio_config=$tmp.fio
fio_out=$tmp.fio.out
cat >$prep_fio_config <<EOF
[randwrite]
name=filler
directory=${subv}
rw=randwrite
nrfiles=${nr_fill}
filesize=${fill_sz}
EOF
_require_fio $fio_config

get_qgroup_usage()
{
	local qgroupid=$1
	local output

	output=$($BTRFS_UTIL_PROG qgroup show --sync --raw $SCRATCH_MNT | \
		 grep "$qgroupid" | $AWK_PROG '{print $3}')
	# The qgroup is auto-removed, this can only happen if its numbers are
	# already all zeros, so here we only need to explicitly echo "0".
	if [ -z "$output" ]; then
		echo "0"
	else
		echo "$output"
	fi
}

get_subvol_usage()
{
	local subvolid=$1

	get_qgroup_usage "0/$subvolid"
}

count_subvol_owned_metadata()
{
	local subvolid=$1

	# We need to sync so that the metadata extents are on disk and visible
	# to dump-tree.
	sync
	# Find nodes and leaves owned by the subvol, then get unique offsets
	# to account for snapshots sharing metadata.
	count=$($BTRFS_UTIL_PROG inspect-internal dump-tree $SCRATCH_DEV | \
		grep "owner $subvolid" | $AWK_PROG '{print $2}' | sort | \
								uniq | wc -l)
	# Output bytes rather than number of metadata blocks.
	echo $(($count * $nodesize))
}

check_qgroup_usage()
{
	local qgroupid=$1
	local expected=$2
	local actual=$(get_qgroup_usage $qgroupid)

	[ $expected -eq $actual ] || \
		echo "qgroup $qgroupid mismatched usage $actual vs $expected"
}

check_subvol_usage()
{
	local subvolid=$1
	local expected_data=$2
	local expected_meta=$(count_subvol_owned_metadata $subvolid)
	local actual=$(get_subvol_usage $subvolid)
	local expected=$(($expected_data + $expected_meta))

	[ $expected -eq $actual ] || \
		echo "subvol $subvolid mismatched usage $actual vs $expected (expected data $expected_data expected meta $expected_meta diff $(($actual - $expected)))"
}

set_subvol_limit()
{
	local subvolid=$1
	local limit=$2

	$BTRFS_UTIL_PROG qgroup limit $2 0/$1 $SCRATCH_MNT
}

# We need the cleaner thread to run to actually delete extents for test
# cases that care about that. The remount wakes up the cleaner thread and sets
# the commit interval to 1s, so the 1.5s sleep is enough to wait for the cleaner
# thread to run.
trigger_cleaner()
{
	_scratch_remount commit=1
	sleep 1.5
}

cycle_mount_check_subvol_usage()
{
	_scratch_cycle_mount
	check_subvol_usage $@
}

do_write()
{
	local file=$1
	local sz=$2

	$XFS_IO_PROG -fc "pwrite -q 0 $sz" $file
}

do_enospc_write()
{
	local file=$1
	local sz=$2

	do_write $file $sz
}

do_falloc()
{
	local file=$1
	local sz=$2

	$XFS_IO_PROG -fc "falloc 0 $sz" $file
}

do_enospc_falloc()
{
	local file=$1
	local sz=$2

	do_falloc $file $sz
}

enable_quota()
{
	local mode=$1

	[ $mode == "n" ] && return
	arg=$([ $mode == "s" ] && echo "--simple")

	$BTRFS_UTIL_PROG quota enable $arg $SCRATCH_MNT
}

get_subvid()
{
	_btrfs_get_subvolid $SCRATCH_MNT subv
}

get_snapid()
{
	_btrfs_get_subvolid $SCRATCH_MNT snap
}

get_nestedid()
{
	_btrfs_get_subvolid $SCRATCH_MNT subv/nested
}

prepare()
{
	_scratch_mkfs >> $seqres.full
	_scratch_mount
	enable_quota "s"
	$BTRFS_UTIL_PROG subvolume create $subv >> $seqres.full
	local subvid=$(get_subvid)
	set_subvol_limit $subvid $limit
	check_subvol_usage $subvid 0

	# Create a bunch of little filler files to generate several levels in
	# the btree, to make snapshotting sharing scenarios complex enough.
	$FIO_PROG $prep_fio_config --output=$fio_out
	check_subvol_usage $subvid $total_fill

	# Create a single file whose extents we will explicitly share/unshare.
	do_write $subv/f $ext_sz
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
}

prepare_snapshotted()
{
	prepare
	$BTRFS_UTIL_PROG subvolume snapshot $subv $snap >> $seqres.full
	check_subvol_usage $(get_subvid) $(($total_fill + $ext_sz))
	check_subvol_usage $(get_snapid) 0
}

prepare_nested()
{
	prepare
	local subvid=$(get_subvid)
	$BTRFS_UTIL_PROG qgroup create 1/100 $SCRATCH_MNT
	$BTRFS_UTIL_PROG qgroup limit $limit 1/100 $SCRATCH_MNT
	$BTRFS_UTIL_PROG qgroup assign 0/$subvid 1/100 $SCRATCH_MNT >> $seqres.full
	$BTRFS_UTIL_PROG subvolume create $nested >> $seqres.full
	local nestedid=$(get_nestedid)
	do_write $nested/f $ext_sz
	check_subvol_usage $nestedid $ext_sz
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	local subv_usage=$(get_subvol_usage $subvid)
	local nested_usage=$(get_subvol_usage $nestedid)
	check_qgroup_usage 1/100 $(($subv_usage + $nested_usage))
}

# Write in a single subvolume, including going over the limit.
basic_accounting()
{
	echo "basic accounting"
	prepare
	local subvid=$(get_subvid)
	rm $subv/f
	check_subvol_usage $subvid $total_fill
	cycle_mount_check_subvol_usage $subvid $total_fill
	do_write $subv/tmp 512M
	rm $subv/tmp
	do_write $subv/tmp 512M
	rm $subv/tmp
	do_enospc_falloc $subv/large_falloc 2G
	do_enospc_write $subv/large 2G
	_scratch_unmount
}

# Write to the same range of a file a bunch of times in a row
# to test extent aware reservations.
reservation_accounting()
{
	echo "reservation accounting"
	prepare
	for i in $(seq 10); do
		do_write $subv/tmp 512M
		rm $subv/tmp
	done
	do_enospc_write $subv/large 2G
	_scratch_unmount
}

# Write in a snapshot.
snapshot_accounting()
{
	echo "snapshot accounting"
	prepare_snapshotted
	local subvid=$(get_subvid)
	local snapid=$(get_snapid)
	touch $snap/f
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	do_write $snap/f $ext_sz
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid $ext_sz
	rm $snap/f
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	rm $subv/f
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	cycle_mount_check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	_scratch_unmount
}

# Delete the original ref first after a snapshot.
delete_snapshot_src_ref()
{
	echo "delete src ref first"
	prepare_snapshotted
	local subvid=$(get_subvid)
	local snapid=$(get_snapid)
	rm $subv/f
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	rm $snap/f
	trigger_cleaner
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	cycle_mount_check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	_scratch_unmount
}

# Delete the snapshot ref first after a snapshot.
delete_snapshot_ref()
{
	echo "delete snapshot ref first"
	prepare_snapshotted
	local subvid=$(get_subvid)
	local snapid=$(get_snapid)
	rm $snap/f
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	rm $subv/f
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	cycle_mount_check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	_scratch_unmount
}

# Delete the snapshotted subvolume after a snapshot.
delete_snapshot_src()
{
	echo "delete snapshot src first"
	prepare_snapshotted
	local subvid=$(get_subvid)
	local snapid=$(get_snapid)
	$BTRFS_UTIL_PROG subvolume delete $subv >> $seqres.full
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	rm $snap/f
	trigger_cleaner
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $snapid 0
	$BTRFS_UTIL_PROG subvolume delete $snap >> $seqres.full
	trigger_cleaner
	check_subvol_usage $subvid 0
	check_subvol_usage $snapid 0
	cycle_mount_check_subvol_usage $subvid 0
	check_subvol_usage $snapid 0
	_scratch_unmount
}

# Delete the snapshot subvolume after a snapshot.
delete_snapshot()
{
	echo "delete snapshot first"
	prepare_snapshotted
	local subvid=$(get_subvid)
	local snapid=$(get_snapid)
	$BTRFS_UTIL_PROG subvolume delete $snap >> $seqres.full
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	check_subvol_usage $snapid 0
	$BTRFS_UTIL_PROG subvolume delete $subv >> $seqres.full
	trigger_cleaner
	check_subvol_usage $subvid 0
	check_subvol_usage $snapid 0
	_scratch_unmount
}

# Write to a subvolume nested in another subvolume.
# Exercises the auto-inheritance feature of simple quotas.
nested_accounting()
{
	echo "nested accounting"
	prepare_nested
	local subvid=$(get_subvid)
	local nestedid=$(get_nestedid)
	rm $subv/f
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $nestedid $ext_sz
	local subv_usage=$(get_subvol_usage $subvid)
	local nested_usage=$(get_subvol_usage $nestedid)
	check_qgroup_usage 1/100 $(($subv_usage + $nested_usage))
	rm $nested/f
	check_subvol_usage $subvid $total_fill
	check_subvol_usage $nestedid 0
	subv_usage=$(get_subvol_usage $subvid)
	nested_usage=$(get_subvol_usage $nestedid)
	check_qgroup_usage 1/100 $(($subv_usage + $nested_usage))
	do_enospc_falloc $nested/large_falloc 2G
	do_enospc_write $nested/large 2G
	_scratch_unmount
}

# Enable simple quotas on a filesystem with existing extents.
enable_mature()
{
	echo "enable mature"
	_scratch_mkfs >> $seqres.full
	_scratch_mount
	$BTRFS_UTIL_PROG subvolume create $subv >> $seqres.full
	local subvid=$(get_subvid)
	do_write $subv/f $ext_sz
	# Sync before enabling squotas to reliably *not* count the writes
	# we did before enabling.
	sync
	enable_quota "s"
	set_subvol_limit $subvid $limit
	_scratch_cycle_mount
	usage=$(get_subvol_usage $subvid)
	[ $usage -lt $ext_sz ] || \
		echo "captured usage from before enable $usage >= $ext_sz"
	do_write $subv/g $ext_sz
	usage=$(get_subvol_usage $subvid)
	[ $usage -lt $ext_sz ] && \
		echo "failed to capture usage after enable $usage < $ext_sz"
	check_subvol_usage $subvid $ext_sz
	rm $subv/f
	check_subvol_usage $subvid $ext_sz
	_scratch_cycle_mount
	rm $subv/g
	check_subvol_usage $subvid 0
	_scratch_unmount
}

# Reflink a file within the subvolume.
reflink_accounting()
{
	echo "reflink"
	prepare
	local subvid=$(get_subvid)
	# Do enough reflinks to prove that they're free. If they counted, then
	# this wouldn't fit in the limit.
	for i in $(seq $(($limit_nr * 2))); do
		_cp_reflink $subv/f $subv/f.i
	done
	# Confirm that there is no additional data usage from the reflinks.
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	_scratch_unmount
}

# Delete the src ref of a reflink first.
delete_reflink_src_ref()
{
	echo "delete reflink src ref"
	prepare
	local subvid=$(get_subvid)
	_cp_reflink $subv/f $subv/f.link
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	rm $subv/f
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	rm $subv/f.link
	check_subvol_usage $subvid $(($total_fill))
	_scratch_unmount
}

# Delete the link ref of a reflink first.
delete_reflink_ref()
{
	echo "delete reflink ref"
	prepare
	local subvid=$(get_subvid)
	_cp_reflink $subv/f $subv/f.link
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	rm $subv/f.link
	check_subvol_usage $subvid $(($total_fill + $ext_sz))
	rm $subv/f
	check_subvol_usage $subvid $(($total_fill))
	_scratch_unmount
}

basic_accounting
reservation_accounting
snapshot_accounting
delete_snapshot_src_ref
delete_snapshot_ref
delete_snapshot_src
delete_snapshot
nested_accounting
enable_mature
reflink_accounting
delete_reflink_src_ref
delete_reflink_ref

# success, all done
status=0
exit