#!/usr/bin/bash
#
# Takes a list of modules and unloads them and all dependent modules.
# If a module cannot be unloaded (e.g. it's in use), an error is returned.
###############################################################################

SCRIPT_NAME="$(basename "$0")"
LCTL=${LCTL:-lctl}
# Seconds to wait for OBD devices to drain before unloading, and to retry a
# failed module unload against its refcount. umount may return before OBD
# devices are fully unconfigured (mgc, lwp/osp imports, obd_zombie), and an
# empty device list does not itself mean a module's refcount has reached
# zero -- there is a short tail while mdt/osd/ptlrpc release their last
# reference. $WAIT_REMAINING is a shared budget for the device drain and
# the module unload retries, decremented only for seconds actually spent
# sleeping in those two phases -- other work in between (kmemleak scan,
# debug dump, lnet unconfigure) does not consume it.
WAIT=${LUSTRE_RMMOD_WAIT:-60}
if [[ -z "$DEBUG" ]]; then
	if [[ -n "$DEBUG_RMMOD" ]]; then
		DEBUG=true
	else
		DEBUG=false
	fi
fi

# Print help message
print_usage() {
	echo "$SCRIPT_NAME -h|--help"
	echo "$SCRIPT_NAME [-d|--debug-kernel] [-w|--wait SECS] [MODULENAME...]"
	echo
	echo -e "\t-d, --debug-kernel\tDisplay lustre kernel debug messages"
	echo -e "\t-h, --help\t\tDisplay this help message"
	echo -e "\t-w, --wait SECS\t\twait for device cleanup and module"
	echo -e "\t\t\t\trefcounts to drain (default 60s)"
	echo -e "\t\t\t\t0 unloads (and potentially fails) immediately"
	echo -e "\tMODULENAME\t\tList of lustre modules to unload."
	echo -e "\t\t\t\tBy default all modules are unloaded."
}

# Wait for all OBD devices to drain and clean up before unloading modules.
#
# No output or an error from "lctl dl" means no devices remain (non-fatal).
# Consumes from the shared $WAIT_REMAINING budget, decremented once per
# second actually spent sleeping here; a slow device drain leaves less
# budget for the module unload retries that follow.
wait_for_devices() {
	local print

	(( WAIT_REMAINING > 0 )) || return 0

	print=$((SECONDS + 20))
	while (( WAIT_REMAINING > 0 )); do
		local num=$($LCTL dl 2>/dev/null | wc -l)
		(( num == 0 )) && return 0

		if (( SECONDS >= print )); then
			echo "$SCRIPT_NAME: wait ${WAIT_REMAINING}s for $num OBD devices" >&2
			(( print+=60 ))
		fi
		sleep 1
		(( WAIT_REMAINING-- )) || true
	done

	if [[ -n "$($LCTL dl 2>/dev/null)" ]]; then
		echo "$SCRIPT_NAME: OBD devices after ${WAIT}s:" >&2
		$LCTL dl >&2
	fi
	return 0
}

# rmmod with retries against the shared $WAIT_REMAINING budget, decremented
# once per second actually spent sleeping here (and in wait_for_devices());
# work in between -- the kmemleak scan, the debug dump, lnet unconfigure --
# does not consume it. An empty "lctl dl" does not imply the module's
# refcnt is zero: there is a short tail after the last OBD device
# disappears where mdt/osd/ptlrpc still hold a module reference, so an
# immediate rmmod can fail "module in use" on a fast teardown (seen on
# single-MDT hosts). Between attempts poll /sys/module/<mod>/refcnt and
# only reattempt once it reads 0 (a module that vanished concurrently
# counts as unloaded). Once the budget is spent, make a final visible
# attempt so a genuinely stuck module still fails with the usual error.
rmmod_wait() {
	local mod=$1
	local refcnt=/sys/module/$mod/refcnt
	local print cnt

	if (( WAIT_REMAINING > 0 )); then
		rmmod $mod 2>/dev/null && return 0

		print=$((SECONDS + 20))
		while (( WAIT_REMAINING > 0 )); do
			sleep 1
			(( WAIT_REMAINING-- )) || true
			[[ -e $refcnt ]] || return 0

			read -r cnt 2>/dev/null < "$refcnt" || cnt=0
			if (( cnt > 0 )); then
				if (( SECONDS >= print )); then
					echo "$SCRIPT_NAME: wait ${WAIT_REMAINING}s for $mod (refcnt $cnt)" >&2
					(( print+=60 ))
				fi
				continue
			fi
			rmmod $mod 2>/dev/null && return 0
		done
	fi

	rmmod $mod
}

# Print kernel debug message for lustre modules
print_debug() {
	local debug_file

	$LCTL mark "$SCRIPT_NAME : Stop debug"
	if [[ $DEBUG_RMMOD == "-" ]]; then
		debug_file="" # dump to stdout
	elif [[ "${DEBUG_RMMOD:0:1}" == "/" ]]; then
		debug_file="$DEBUG_RMMOD"
	else
		debug_file=$TMP/${DEBUG_RMMOD:-debug}
	fi
	echo "Dump memory leak logs to $debug_file"
	$LCTL debug_kernel $debug_file
	DEBUG=false
}

# Unload all modules dependent on $1 (exclude removal of $1)
unload_dep_modules_exclusive() {
	local MODULE=$1

	local DEPS="$(lsmod | awk '($1 == "'$MODULE'") { print $4 }')"
	for SUBMOD in $(echo $DEPS | tr ',' ' '); do
		unload_dep_modules_inclusive $SUBMOD || return 1
	done
	return 0
}

# Unload all modules dependent on $1 (include removal of $1)
unload_dep_modules_inclusive() {
	local MODULE=$1

	# if $MODULE not loaded, return 0
	lsmod | grep -E -q "^\<$MODULE\>" || return 0
	unload_dep_modules_exclusive $MODULE || return 1

	if $DEBUG; then
		if [ "$MODULE" = 'libcfs' ]; then
			print_debug
		fi
		$LCTL mark "$SCRIPT_NAME : Unload $MODULE"
	fi

	rmmod_wait $MODULE || return 1
	return 0
}

declare -a modules
while (( $# > 0 )); do
	case "$1" in
		-d|--debug-kernel)
			if lsmod | grep -E -q '^libcfs'; then
				DEBUG='true'
			else
				echo "Debug unavailable: libcfs not loaded" >&2
			fi
			;;
		-h|--help)
			print_usage >&2
			exit 0
			;;
		-w|--wait)
			shift
			WAIT="$1"
			if ! [[ "$WAIT" =~ ^[0-9]+$ ]]; then
				echo "Error: --wait needs number of seconds" >&2
				print_usage >&2
				exit 2
			fi
			;;
		-*)
			echo "Error invalid option '$1'" >&2
			print_usage >&2
			exit 2
			;;
		*)
			modules+=("$1")
			;;
	esac
	shift
done

# To maintain backwards compatibility, ldiskfs and libcfs must be
# unloaded if no parameters are given, or if only the ldiskfs parameter
# is given. It's ugly, but is needed to emulate the prior functionality
if (( ${#modules[@]} == 0 )) || [[ "${modules[*]}" == "ldiskfs" ]]; then
	unload_all=true
	modules=('lnet_selftest' 'ldiskfs' 'libcfs')
else
	unload_all=false
fi

# Shared budget for the device drain and per-module refcnt retries.
WAIT_REMAINING=$WAIT

wait_for_devices

export KMEMLEAK=${KMEMLEAK:-/sys/kernel/debug/kmemleak}
KMEMLEAK_MODS=/tmp/kmemleak-modules-list.txt
if [[ -w $KMEMLEAK ]]; then
	if ! echo scan > $KMEMLEAK 2>&1; then
		echo "kmemleak disabled"
		export KMEMLEAK=disabled
	else
		kmemleak_pre=/tmp/kmemleak-pre-unload.txt

		cat /proc/modules > $KMEMLEAK_MODS
		cat $KMEMLEAK > $kmemleak_pre
		[[ -s $kmemleak_pre ]] && logger -t leak-pre -f $kmemleak_pre
		rm -f $kmemleak_pre
		# Clear everything here so that only new leaks show up
		# after module unload
		echo clear > $KMEMLEAK
	fi
fi

# Manage debug
if $DEBUG; then
	echo "Lustre debug parameters:" >&2
	$LCTL get_param debug >&2
	$LCTL get_param debug_mb >&2

	$LCTL mark "$SCRIPT_NAME : Start debug"
fi

if $unload_all; then
	unload_dep_modules_inclusive 'ptlrpc' || exit 1
	# LNet may have an internal ref which can prevent LND modules from
	# unloading. Try to drop it before unloading modules.
	# NB: we squelch stderr because lnetctl/lctl may complain about
	# LNet being "busy", but this is normal. We're making a best effort
	# here.
	# Prefer lnetctl if it is present
	if [ -n "$(which lnetctl 2>/dev/null)" ]; then
		lnetctl lnet unconfigure 2>/dev/null
	elif [ -n "$(which lctl 2>/dev/null)" ]; then
		lctl net down 2>/dev/null | grep -v "LNET ready to unload"
	fi
fi

for mod in ${modules[*]}; do
	unload_dep_modules_inclusive $mod || exit 1
done

if $DEBUG; then
	print_debug
fi

if [[ -f $KMEMLEAK ]]; then
	kmemleak_post=/tmp/kmemleak-post-unload.txt

	echo scan > $KMEMLEAK 2>&1 | grep -v "Device or resource busy"
	cat $KMEMLEAK > $kmemleak_post
	[[ -s $kmemleak_post ]] && logger -t leak-mods -f $KMEMLEAK_MODS &&
		logger -t leak-post -f $kmemleak_post
	rm -f $kmemleak_post $KMEMLEAK_MODS
fi

exit 0
